diff --git a/.github/workflows/VeniceCI-CompatibilityTests.yml b/.github/workflows/VeniceCI-CompatibilityTests.yml index 058cdf71144..3dd0a78ca9d 100644 --- a/.github/workflows/VeniceCI-CompatibilityTests.yml +++ b/.github/workflows/VeniceCI-CompatibilityTests.yml @@ -1,4 +1,4 @@ -# GitHub Actions workflow for running compatibility tests: Avro, Alpini unit, Alpini functional tests, and Pulsar Venice integration tests +# GitHub Actions workflow for running compatibility tests: Avro and Pulsar Venice integration tests name: TestsDeCompatibilite @@ -50,65 +50,6 @@ jobs: path: ${{ github.job }}-jdk${{ matrix.jdk }}-logs.tar.gz retention-days: 30 - AlpiniUnitTests: - strategy: - fail-fast: false - matrix: - jdk: [8, 11, 17] - runs-on: ubuntu-latest - timeout-minutes: 120 - steps: - - uses: actions/checkout@v4 - with: - # Checkout as many commits as needed for the diff - fetch-depth: 2 - - uses: actions/checkout@v4 - with: - fetch-depth: 0 - - name: Set up JDK - uses: actions/setup-java@v4 - with: - java-version: ${{ matrix.jdk }} - distribution: 'temurin' - cache: 'gradle' - # - name: Allow Deprecated TLS versions for Alpini tests - # run: | - # echo "java.security file before modifications: " - # cat "$JAVA_HOME/conf/security/java.security" - - # # This is possibly flaky but - # sed -i 's/TLSv1, //g' "$JAVA_HOME/conf/security/java.security" # Allow TLSv1 - # sed -i 's/TLSv1.1, //g' "$JAVA_HOME/conf/security/java.security" # Allow TLSv1.1 - - # echo "java.security file after modifications: " - # cat "$JAVA_HOME/conf/security/java.security" - - shell: bash - run: | - git remote set-head origin --auto - git remote add upstream https://github.com/linkedin/venice - git fetch upstream - - name: Setup Gradle - uses: gradle/actions/setup-gradle@v4 - with: - add-job-summary: never - - name: Run alpini unit tests - run: ./gradlew --continue --no-daemon -DmaxParallelForks=1 alpiniUnitTest - - name: Package Build Artifacts - if: (success() || failure()) - shell: bash - run: | - mkdir ${{ github.job }}-artifacts - find . -path "**/build/reports/*" -or -path "**/build/test-results/*" > artifacts.list - rsync -R --files-from=artifacts.list . ${{ github.job }}-artifacts - tar -zcvf ${{ github.job }}-jdk${{ matrix.jdk }}-logs.tar.gz ${{ github.job }}-artifacts - - name: Upload Build Artifacts - if: (success() || failure()) - uses: actions/upload-artifact@v4 - with: - name: ${{ github.job }}-jdk${{ matrix.jdk }} - path: ${{ github.job }}-jdk${{ matrix.jdk }}-logs.tar.gz - retention-days: 30 - PulsarVeniceIntegrationTests: strategy: fail-fast: false @@ -178,7 +119,7 @@ jobs: strategy: fail-fast: false runs-on: ubuntu-latest - needs: [AvroCompatibilityTests, AlpiniUnitTests, PulsarVeniceIntegrationTests] + needs: [AvroCompatibilityTests, PulsarVeniceIntegrationTests] timeout-minutes: 120 steps: - name: AllIsWell diff --git a/.github/workflows/VeniceCI-E2ETests.yml b/.github/workflows/VeniceCI-E2ETests.yml index 1fabe472a35..15c77e90ecf 100644 --- a/.github/workflows/VeniceCI-E2ETests.yml +++ b/.github/workflows/VeniceCI-E2ETests.yml @@ -108,6 +108,101 @@ jobs: key: ${{ secrets.BUILDPULSE_ACCESS_KEY_ID }} secret: ${{ secrets.BUILDPULSE_SECRET_ACCESS_KEY }} + IntegrationTests_1001: + name: IntegrationTests_1001 + strategy: + fail-fast: false + matrix: + jdk: [17] + runs-on: ubuntu-latest + permissions: + id-token: write + contents: read + checks: write + pull-requests: write + issues: write + timeout-minutes: 120 + concurrency: + group: ${{ github.workflow }}-${{ github.event.pull_request.number || github.ref }}-jdk${{ matrix.jdk }}-IntegrationTests_1001 + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + steps: + - uses: actions/checkout@v4 + with: + fetch-depth: 0 + - name: Set up JDK + uses: actions/setup-java@v4 + with: + java-version: ${{ matrix.jdk }} + distribution: 'temurin' + cache: 'gradle' + - shell: bash + run: | + git remote set-head origin --auto + git remote add upstream https://github.com/linkedin/venice + git fetch upstream + - name: Setup Gradle + uses: gradle/actions/setup-gradle@v4 + with: + add-job-summary: never + - name: Run Integration Tests + run: ./gradlew --continue --no-daemon -DforkEvery=1 -DmaxParallelForks=1 integrationTests_1001 + - name: Package Build Artifacts + if: success() || failure() + shell: bash + run: | + mkdir ${{ github.job }}-artifacts + echo "Repository owner: ${{ github.repository_owner }}" + echo "Repository name: ${{ github.repository }}" + echo "event name: ${{ github.event_name }}" + find . -path "**/build/reports/*" -or -path "**/build/test-results/*" > artifacts.list + rsync -R --files-from=artifacts.list . ${{ github.job }}-artifacts + tar -zcvf ${{ github.job }}-jdk${{ matrix.jdk }}-logs.tar.gz ${{ github.job }}-artifacts + - name: Generate Fork Repo Test Reports + if: ${{ (github.repository_owner != 'linkedin') && (success() || failure()) }} + uses: dorny/test-reporter@v1.9.1 + env: + NODE_OPTIONS: --max-old-space-size=9182 + with: + token: ${{ secrets.GITHUB_TOKEN }} + name: ${{ github.job }} Test Reports # Name where it report the test results + path: '**/TEST-*.xml' + fail-on-error: 'false' + max-annotations: '10' + list-tests: 'all' + list-suites: 'all' + reporter: java-junit + - name: Publish Test Report + continue-on-error: true + env: + NODE_OPTIONS: "--max_old_space_size=8192" + uses: mikepenz/action-junit-report@v5 + if: always() + with: + check_name: ${{ github.job }}-jdk${{ matrix.jdk }} Report + comment: false + annotate_only: true + flaky_summary: true + commit: ${{github.event.workflow_run.head_sha}} + detailed_summary: true + report_paths: '**/build/test-results/test/TEST-*.xml' + - name: Upload Build Artifacts + if: success() || failure() + uses: actions/upload-artifact@v4 + with: + name: ${{ github.job }} + path: ${{ github.job }}-jdk${{ matrix.jdk }}-logs.tar.gz + retention-days: 30 + - name: Upload test results to BuildPulse for flaky test detection + if: github.event_name == 'push' && github.ref == 'refs/heads/main' && !cancelled() + uses: buildpulse/buildpulse-action@main + with: + account: 100582612927 + repository: 100441445875 + path: | + **/TEST-*.xml + key: ${{ secrets.BUILDPULSE_ACCESS_KEY_ID }} + secret: ${{ secrets.BUILDPULSE_SECRET_ACCESS_KEY }} + IntegrationTests_1010: name: IntegrationTests_1010 strategy: @@ -2964,7 +3059,7 @@ jobs: matrix: jdk: [17] runs-on: ubuntu-latest - needs: [IntegrationTests_1000, IntegrationTests_1010, IntegrationTests_1020, IntegrationTests_1030, IntegrationTests_1040, IntegrationTests_1050, IntegrationTests_1060, IntegrationTests_1070, IntegrationTests_1080, IntegrationTests_1090, IntegrationTests_1100, IntegrationTests_1110, IntegrationTests_1120, IntegrationTests_1130, IntegrationTests_1200, IntegrationTests_1210, IntegrationTests_1220, IntegrationTests_1230, IntegrationTests_1240, IntegrationTests_1250, IntegrationTests_1260, IntegrationTests_1270, IntegrationTests_1280, IntegrationTests_1400, IntegrationTests_1410, IntegrationTests_1420, IntegrationTests_1430, IntegrationTests_1440, IntegrationTests_1500, IntegrationTests_1550, IntegrationTests_9999] + needs: [IntegrationTests_1000, IntegrationTests_1001, IntegrationTests_1010, IntegrationTests_1020, IntegrationTests_1030, IntegrationTests_1040, IntegrationTests_1050, IntegrationTests_1060, IntegrationTests_1070, IntegrationTests_1080, IntegrationTests_1090, IntegrationTests_1100, IntegrationTests_1110, IntegrationTests_1120, IntegrationTests_1130, IntegrationTests_1200, IntegrationTests_1210, IntegrationTests_1220, IntegrationTests_1230, IntegrationTests_1240, IntegrationTests_1250, IntegrationTests_1260, IntegrationTests_1270, IntegrationTests_1280, IntegrationTests_1400, IntegrationTests_1410, IntegrationTests_1420, IntegrationTests_1430, IntegrationTests_1440, IntegrationTests_1500, IntegrationTests_1550, IntegrationTests_9999] timeout-minutes: 20 steps: - name: NoOp diff --git a/.github/workflows/VeniceCI-StaticAnalysisAndUnitTests.yml b/.github/workflows/VeniceCI-StaticAnalysisAndUnitTests.yml index 5c26ed0e043..44f63d309bc 100644 --- a/.github/workflows/VeniceCI-StaticAnalysisAndUnitTests.yml +++ b/.github/workflows/VeniceCI-StaticAnalysisAndUnitTests.yml @@ -67,46 +67,66 @@ jobs: uses: ./.github/workflows/UnitTests-core.yml with: artifact_suffix: clients - arg: :clients:venice-admin-tool:jacocoTestCoverageVerification :clients:venice-admin-tool:diffCoverage - :clients:venice-producer:jacocoTestCoverageVerification :clients:venice-producer:diffCoverage - :integrations:venice-pulsar:jacocoTestCoverageVerification :integrations:venice-pulsar:diffCoverage - :clients:venice-client:jacocoTestCoverageVerification :clients:venice-client:diffCoverage - :clients:venice-push-job:jacocoTestCoverageVerification :clients:venice-push-job:diffCoverage - :integrations:venice-samza:jacocoTestCoverageVerification :integrations:venice-samza:diffCoverage - :clients:venice-thin-client:jacocoTestCoverageVerification :clients:venice-thin-client:diffCoverage --continue + arg: + # Cannot use :clients:recursiveDiffCoverage because that would include DVC, and we want this one to run in server... + :clients:venice-admin-tool:diffCoverage + :clients:venice-producer:diffCoverage + :clients:venice-client:diffCoverage + :clients:venice-push-job:diffCoverage + :clients:venice-thin-client:diffCoverage + --continue + + Integrations: + uses: ./.github/workflows/UnitTests-core.yml + with: + artifact_suffix: integrations + arg: + :integrations:recursiveDiffCoverage + --continue Internal: uses: ./.github/workflows/UnitTests-core.yml with: artifact_suffix: internal - arg: :internal:venice-client-common:jacocoTestCoverageVerification :internal:venice-client-common:diffCoverage - :internal:venice-common:jacocoTestCoverageVerification :internal:venice-common:diffCoverage - :internal:venice-jdk-compatibility-test:jacocoTestCoverageVerification :internal:venice-jdk-compatibility-test:diffCoverage - :internal:venice-test-common:jacocoTestCoverageVerification :internal:venice-test-common:diffCoverage --continue - + arg: + # Cannot use :internal:recursiveDiffCoverage because that would include the avro compat test, and we want this one to run in the Compatibility group (TODO: move it out of internal?)... + :internal:venice-client-common:diffCoverage + :internal:venice-common:diffCoverage + :internal:venice-jdk-compatibility-test:diffCoverage + :internal:venice-test-common:diffCoverage + --continue Controller: uses: ./.github/workflows/UnitTests-core.yml with: artifact_suffix: controller - arg: :services:venice-controller:jacocoTestCoverageVerification :services:venice-controller:diffCoverage --continue + arg: + :services:venice-controller:diffCoverage + --continue + Server: uses: ./.github/workflows/UnitTests-core.yml with: artifact_suffix: server - arg: :clients:da-vinci-client:jacocoTestCoverageVerification :clients:da-vinci-client:diffCoverage - :services:venice-server:jacocoTestCoverageVerification :services:venice-server:diffCoverage --continue + arg: + :clients:da-vinci-client:diffCoverage + :services:venice-server:diffCoverage + --continue + Router: uses: ./.github/workflows/UnitTests-core.yml with: artifact_suffix: router - arg: :services:venice-router:jacocoTestCoverageVerification :services:venice-router:diffCoverage --continue + arg: + :services:venice-router:diffCoverage + alpiniUnitTest + --continue StaticAnalysisAndUnitTestsCompletionCheck: strategy: fail-fast: false runs-on: ubuntu-latest - needs: [ValidateGradleWrapper, StaticAnalysis, Clients, Internal, Controller, Server, Router] + needs: [ValidateGradleWrapper, StaticAnalysis, Clients, Integrations, Internal, Controller, Server, Router] timeout-minutes: 120 if: success() || failure() # Always run this job, regardless of previous job status steps: @@ -128,6 +148,10 @@ jobs: echo "Internal module unit tests failed." exit 1 fi + if [ "${{ needs.Integrations.result }}" != "success" ]; then + echo "Integrations module unit tests failed." + exit 1 + fi if [ "${{ needs.Controller.result }}" != "success" ]; then echo "Controller module unit tests failed." exit 1 diff --git a/.gitignore b/.gitignore index 992fbb6e907..44b4c8bcd52 100644 --- a/.gitignore +++ b/.gitignore @@ -24,3 +24,4 @@ Gemfile.lock .bundles_cache docs/vendor/ clients/da-vinci-client/classHash*.txt +integrations/venice-duckdb/classHash*.txt diff --git a/build.gradle b/build.gradle index 38136cbc143..ac76c4e771d 100644 --- a/build.gradle +++ b/build.gradle @@ -10,9 +10,9 @@ plugins { id 'maven-publish' id 'com.diffplug.spotless' version '6.12.0' id 'com.dorongold.task-tree' version '2.1.0' - id 'com.github.johnrengelman.shadow' version '6.1.0' apply false + id 'com.github.johnrengelman.shadow' version '7.1.2' apply false id 'com.github.spotbugs' version '4.8.0' apply false - id 'org.gradle.test-retry' version '1.5.0' apply false + id 'org.gradle.test-retry' version '1.6.0' apply false id 'com.form.diff-coverage' version '0.9.5' apply false id 'me.champeau.jmh' version '0.6.7' apply false id 'io.github.lhotari.gradle-nar-plugin' version '0.5.1' apply false @@ -82,12 +82,14 @@ ext.libraries = [ commonsLang: 'commons-lang:commons-lang:2.6', conscrypt: 'org.conscrypt:conscrypt-openjdk-uber:2.5.2', d2: "com.linkedin.pegasus:d2:${pegasusVersion}", + duckdbJdbc: "org.duckdb:duckdb_jdbc:1.2.0-20250116.012809-118", // TODO: Remove SNAPSHOT when the real release is published! failsafe: 'net.jodah:failsafe:2.4.0', fastUtil: 'it.unimi.dsi:fastutil:8.3.0', grpcNettyShaded: "io.grpc:grpc-netty-shaded:${grpcVersion}", grpcProtobuf: "io.grpc:grpc-protobuf:${grpcVersion}", grpcServices: "io.grpc:grpc-services:${grpcVersion}", grpcStub: "io.grpc:grpc-stub:${grpcVersion}", + grpcTesting: "io.grpc:grpc-testing:${grpcVersion}", hadoopCommon: "org.apache.hadoop:hadoop-common:${hadoopVersion}", hadoopHdfs: "org.apache.hadoop:hadoop-hdfs:${hadoopVersion}", httpAsyncClient: 'org.apache.httpcomponents:httpasyncclient:4.1.5', @@ -116,6 +118,11 @@ ext.libraries = [ mapreduceClientJobClient: "org.apache.hadoop:hadoop-mapreduce-client-jobclient:${hadoopVersion}", mockito: 'org.mockito:mockito-core:4.11.0', netty: 'io.netty:netty-all:4.1.74.Final', + opentelemetryApi: "io.opentelemetry:opentelemetry-api:${openTelemetryVersion}", + opentelemetrySdk: "io.opentelemetry:opentelemetry-sdk:${openTelemetryVersion}", + opentelemetryExporterLogging: "io.opentelemetry:opentelemetry-exporter-logging:${openTelemetryVersion}", + opentelemetryExporterOtlp: "io.opentelemetry:opentelemetry-exporter-otlp:${openTelemetryVersion}", + opentelemetryExporterCommon: "io.opentelemetry:opentelemetry-exporter-common:${openTelemetryVersion}", oss: 'org.sonatype.oss:oss-parent:7', pulsarClient: "${pulsarGroup}:pulsar-client:${pulsarVersion}", pulsarIoCore: "${pulsarGroup}:pulsar-io-core:${pulsarVersion}", @@ -142,12 +149,7 @@ ext.libraries = [ xerces: 'xerces:xercesImpl:2.9.1', zkclient: 'com.101tec:zkclient:0.7', // For Kafka AdminUtils zookeeper: 'org.apache.zookeeper:zookeeper:3.6.3', - zstd: 'com.github.luben:zstd-jni:1.5.2-3', - opentelemetryApi: "io.opentelemetry:opentelemetry-api:${openTelemetryVersion}", - opentelemetrySdk: "io.opentelemetry:opentelemetry-sdk:${openTelemetryVersion}", - opentelemetryExporterLogging: "io.opentelemetry:opentelemetry-exporter-logging:${openTelemetryVersion}", - opentelemetryExporterOtlp: "io.opentelemetry:opentelemetry-exporter-otlp:${openTelemetryVersion}", - opentelemetryExporterCommon: "io.opentelemetry:opentelemetry-exporter-common:${openTelemetryVersion}" + zstd: 'com.github.luben:zstd-jni:1.5.2-3' ] group = 'com.linkedin.venice' @@ -188,9 +190,14 @@ subprojects { def isLeafSubModule = project.childProjects.isEmpty() + // We consider a sub-module to be a "proto module" if it has any proto schemas defined under the right path. + def protoDir = new File(project.projectDir, "src/main/proto"); + def isProtoModule = protoDir != null && protoDir.list() != null && protoDir.list().size() != 0 + apply { plugin 'idea' plugin 'java-library' + plugin 'com.form.diff-coverage' plugin 'com.github.spotbugs' plugin 'org.gradle.test-retry' plugin 'org.checkerframework' @@ -199,8 +206,9 @@ subprojects { if (isLeafSubModule) { apply { plugin 'jacoco' - plugin 'com.form.diff-coverage' - plugin 'com.google.protobuf' + if (isProtoModule) { + plugin 'com.google.protobuf' + } } } @@ -218,7 +226,7 @@ subprojects { //withJavadocJar() } - if (isLeafSubModule) { + if (isLeafSubModule && isProtoModule) { protobuf { protoc { artifact = 'com.google.protobuf:protoc:' + protobufVersion @@ -498,10 +506,6 @@ subprojects { xml.enabled = true html.enabled = true } - - doLast { - parseJacocoXml("$buildDir/reports/jacoco/test/jacocoTestReport.xml") - } } afterEvaluate { @@ -517,6 +521,12 @@ subprojects { value = 'COVEREDRATIO' minimum = threshold } + // Ignore generate files + afterEvaluate { + classDirectories.setFrom(files(classDirectories.files.collect { fileTree(dir: it, exclude: [ + '**/com/linkedin/venice/protocols/**', + ])})) + } } } } @@ -582,6 +592,13 @@ subprojects { } } +// 2nd round of subprojects configuration... the 1st round must be fully done for all submodules for this one to work. +subprojects { + task recursiveDiffCoverage { + dependsOn subprojects.diffCoverage + } +} + task aggregateJavadoc(type: Javadoc) { source subprojects.collect { project -> project.sourceSets.main.allJava @@ -829,4 +846,4 @@ task listSubprojects { println "${subproject.name}" } } -} \ No newline at end of file +} diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java index 278fa9dcef5..855007b9229 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java @@ -9,7 +9,8 @@ import com.linkedin.davinci.blobtransfer.BlobTransferManager; import com.linkedin.davinci.blobtransfer.BlobTransferUtil; -import com.linkedin.davinci.client.DaVinciRecordTransformerFunctionalInterface; +import com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferTableFormat; +import com.linkedin.davinci.client.DaVinciRecordTransformerConfig; import com.linkedin.davinci.compression.StorageEngineBackedCompressorFactory; import com.linkedin.davinci.config.StoreBackendConfig; import com.linkedin.davinci.config.VeniceConfigLoader; @@ -121,7 +122,7 @@ public DaVinciBackend( Optional> managedClients, ICProvider icProvider, Optional cacheConfig, - DaVinciRecordTransformerFunctionalInterface recordTransformerFunction) { + DaVinciRecordTransformerConfig recordTransformerConfig) { LOGGER.info("Creating Da Vinci backend with managed clients: {}", managedClients); try { VeniceServerConfig backendConfig = configLoader.getVeniceServerConfig(); @@ -270,7 +271,7 @@ public DaVinciBackend( false, compressorFactory, cacheBackend, - recordTransformerFunction, + recordTransformerConfig, true, // TODO: consider how/if a repair task would be valid for Davinci users? null, @@ -278,6 +279,7 @@ public DaVinciBackend( Optional.empty(), // TODO: It would be good to monitor heartbeats like this from davinci, but needs some work null, + null, null); ingestionService.start(); @@ -293,7 +295,7 @@ public DaVinciBackend( } if (backendConfig.isBlobTransferManagerEnabled()) { - if (recordTransformerFunction != null) { + if (recordTransformerConfig != null) { throw new VeniceException("DaVinciRecordTransformer doesn't support blob transfer."); } @@ -311,7 +313,10 @@ public DaVinciBackend( backendConfig.getMaxConcurrentSnapshotUser(), backendConfig.getSnapshotRetentionTimeInMin(), backendConfig.getBlobTransferMaxTimeoutInMin(), - aggVersionedBlobTransferStats); + aggVersionedBlobTransferStats, + backendConfig.getRocksDBServerConfig().isRocksDBPlainTableFormatEnabled() + ? BlobTransferTableFormat.PLAIN_TABLE + : BlobTransferTableFormat.BLOCK_BASED_TABLE); } else { aggVersionedBlobTransferStats = null; blobTransferManager = null; diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobSnapshotManager.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobSnapshotManager.java index cc6e5929bef..5ee5620ac2d 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobSnapshotManager.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobSnapshotManager.java @@ -63,6 +63,7 @@ public class BlobSnapshotManager { private final StorageMetadataService storageMetadataService; private final int maxConcurrentUsers; private final long snapshotRetentionTimeInMillis; + private final BlobTransferUtils.BlobTransferTableFormat blobTransferTableFormat; private final Lock lock = new ReentrantLock(); /** @@ -73,13 +74,14 @@ public BlobSnapshotManager( StorageEngineRepository storageEngineRepository, StorageMetadataService storageMetadataService, int maxConcurrentUsers, - int snapshotRetentionTimeInMin) { + int snapshotRetentionTimeInMin, + BlobTransferUtils.BlobTransferTableFormat transferTableFormat) { this.readOnlyStoreRepository = readOnlyStoreRepository; this.storageEngineRepository = storageEngineRepository; this.storageMetadataService = storageMetadataService; this.maxConcurrentUsers = maxConcurrentUsers; this.snapshotRetentionTimeInMillis = TimeUnit.MINUTES.toMillis(snapshotRetentionTimeInMin); - + this.blobTransferTableFormat = transferTableFormat; this.concurrentSnapshotUsers = new VeniceConcurrentHashMap<>(); this.snapshotTimestamps = new VeniceConcurrentHashMap<>(); this.snapshotMetadataRecords = new VeniceConcurrentHashMap<>(); @@ -99,7 +101,8 @@ public BlobSnapshotManager( storageEngineRepository, storageMetadataService, DEFAULT_MAX_CONCURRENT_USERS, - DEFAULT_SNAPSHOT_RETENTION_TIME_IN_MIN); + DEFAULT_SNAPSHOT_RETENTION_TIME_IN_MIN, + BlobTransferUtils.BlobTransferTableFormat.BLOCK_BASED_TABLE); } /** @@ -364,4 +367,12 @@ public BlobTransferPartitionMetadata prepareMetadata(BlobTransferPayload blobTra offsetRecordByte, storeVersionStateByte); } + + /** + * Get the current snapshot format, which is a config value. + * @return the transfer table format, BLOCK_BASED_TABLE or PLAIN_TABLE. + */ + public BlobTransferUtils.BlobTransferTableFormat getBlobTransferTableFormat() { + return this.blobTransferTableFormat; + } } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferManager.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferManager.java index d1810e6d87d..e6fc5045184 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferManager.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferManager.java @@ -1,5 +1,6 @@ package com.linkedin.davinci.blobtransfer; +import com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferTableFormat; import com.linkedin.davinci.stats.AggVersionedBlobTransferStats; import com.linkedin.venice.annotation.Experimental; import com.linkedin.venice.exceptions.VenicePeersNotFoundException; @@ -26,13 +27,17 @@ public interface BlobTransferManager extends AutoCloseable { * @param storeName * @param version * @param partition + * @param requestTableFormat the table format defined in config (PLAIN_TABLE or BLOCK_BASED_TABLE). * @return the InputStream of the blob. The return type is experimental and may change in the future. * @throws VenicePeersNotFoundException when the peers are not found for the requested blob. Other exceptions may be * thrown, but it's wrapped inside the CompletionStage. */ @Experimental - CompletionStage get(String storeName, int version, int partition) - throws VenicePeersNotFoundException; + CompletionStage get( + String storeName, + int version, + int partition, + BlobTransferTableFormat requestTableFormat) throws VenicePeersNotFoundException; /** * Put the blob for the given storeName and partition diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferPayload.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferPayload.java index 3de5e572f00..7c11d952082 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferPayload.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferPayload.java @@ -3,6 +3,7 @@ import static com.linkedin.venice.store.rocksdb.RocksDBUtils.composePartitionDbDir; import static com.linkedin.venice.store.rocksdb.RocksDBUtils.composeSnapshotDir; +import com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferTableFormat; import com.linkedin.venice.utils.Utils; @@ -14,12 +15,19 @@ public class BlobTransferPayload { private final String topicName; private final String partitionDir; private final String storeName; - - public BlobTransferPayload(String baseDir, String storeName, int version, int partition) { + private final BlobTransferTableFormat requestTableFormat; + + public BlobTransferPayload( + String baseDir, + String storeName, + int version, + int partition, + BlobTransferTableFormat requestTableFormat) { this.partition = partition; this.storeName = storeName; this.topicName = storeName + "_v" + version; this.partitionDir = composePartitionDbDir(baseDir, topicName, partition); + this.requestTableFormat = requestTableFormat; } public String getPartitionDir() { @@ -45,4 +53,8 @@ public int getPartition() { public String getStoreName() { return storeName; } + + public BlobTransferTableFormat getRequestTableFormat() { + return requestTableFormat; + } } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferUtil.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferUtil.java index f857de5ec34..f3058342e07 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferUtil.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferUtil.java @@ -41,7 +41,8 @@ public static BlobTransferManager getP2PBlobTransferManagerForDVCAndStart( int maxConcurrentSnapshotUser, int snapshotRetentionTimeInMin, int blobTransferMaxTimeoutInMin, - AggVersionedBlobTransferStats aggVersionedBlobTransferStats) { + AggVersionedBlobTransferStats aggVersionedBlobTransferStats, + BlobTransferUtils.BlobTransferTableFormat transferSnapshotTableFormat) { return getP2PBlobTransferManagerForDVCAndStart( p2pTransferPort, p2pTransferPort, @@ -53,7 +54,8 @@ public static BlobTransferManager getP2PBlobTransferManagerForDVCAndStart( maxConcurrentSnapshotUser, snapshotRetentionTimeInMin, blobTransferMaxTimeoutInMin, - aggVersionedBlobTransferStats); + aggVersionedBlobTransferStats, + transferSnapshotTableFormat); } public static BlobTransferManager getP2PBlobTransferManagerForDVCAndStart( @@ -67,14 +69,16 @@ public static BlobTransferManager getP2PBlobTransferManagerForDVCAndStart( int maxConcurrentSnapshotUser, int snapshotRetentionTimeInMin, int blobTransferMaxTimeoutInMin, - AggVersionedBlobTransferStats aggVersionedBlobTransferStats) { + AggVersionedBlobTransferStats aggVersionedBlobTransferStats, + BlobTransferUtils.BlobTransferTableFormat transferSnapshotTableFormat) { try { BlobSnapshotManager blobSnapshotManager = new BlobSnapshotManager( readOnlyStoreRepository, storageEngineRepository, storageMetadataService, maxConcurrentSnapshotUser, - snapshotRetentionTimeInMin); + snapshotRetentionTimeInMin, + transferSnapshotTableFormat); AbstractAvroStoreClient storeClient = new AvroGenericStoreClientImpl<>(getTransportClient(clientConfig), false, clientConfig); BlobTransferManager manager = new NettyP2PBlobTransferManager( @@ -111,14 +115,16 @@ public static BlobTransferManager getP2PBlobTransferManagerForServerAndSta int maxConcurrentSnapshotUser, int snapshotRetentionTimeInMin, int blobTransferMaxTimeoutInMin, - AggVersionedBlobTransferStats aggVersionedBlobTransferStats) { + AggVersionedBlobTransferStats aggVersionedBlobTransferStats, + BlobTransferUtils.BlobTransferTableFormat transferSnapshotTableFormat) { try { BlobSnapshotManager blobSnapshotManager = new BlobSnapshotManager( readOnlyStoreRepository, storageEngineRepository, storageMetadataService, maxConcurrentSnapshotUser, - snapshotRetentionTimeInMin); + snapshotRetentionTimeInMin, + transferSnapshotTableFormat); BlobTransferManager manager = new NettyP2PBlobTransferManager( new P2PBlobTransferService(p2pTransferServerPort, baseDir, blobTransferMaxTimeoutInMin, blobSnapshotManager), new NettyFileTransferClient(p2pTransferClientPort, baseDir, storageMetadataService), diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferUtils.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferUtils.java index 5da87d39eae..3cd3a1f5ad8 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferUtils.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/BlobTransferUtils.java @@ -22,6 +22,10 @@ public enum BlobTransferType { FILE, METADATA } + public enum BlobTransferTableFormat { + PLAIN_TABLE, BLOCK_BASED_TABLE + } + /** * Check if the HttpResponse message is for metadata. * @param msg the HttpResponse message diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/NettyP2PBlobTransferManager.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/NettyP2PBlobTransferManager.java index d31e83c92bf..5f520d5ef10 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/NettyP2PBlobTransferManager.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/NettyP2PBlobTransferManager.java @@ -2,6 +2,7 @@ import static com.linkedin.davinci.blobtransfer.BlobTransferUtils.getThroughputPerPartition; +import com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferTableFormat; import com.linkedin.davinci.blobtransfer.client.NettyFileTransferClient; import com.linkedin.davinci.blobtransfer.server.P2PBlobTransferService; import com.linkedin.davinci.stats.AggVersionedBlobTransferStats; @@ -72,8 +73,11 @@ public void start() throws Exception { } @Override - public CompletionStage get(String storeName, int version, int partition) - throws VenicePeersNotFoundException { + public CompletionStage get( + String storeName, + int version, + int partition, + BlobTransferTableFormat tableFormat) throws VenicePeersNotFoundException { CompletableFuture resultFuture = new CompletableFuture<>(); // 1. Discover peers for the requested blob BlobPeersDiscoveryResponse response = peerFinder.discoverBlobPeers(storeName, version, partition); @@ -92,7 +96,7 @@ public CompletionStage get(String storeName, int version, int parti .info("Discovered peers {} for store {} version {} partition {}", discoverPeers, storeName, version, partition); // 2: Process peers sequentially to fetch the blob - processPeersSequentially(discoverPeers, storeName, version, partition, resultFuture); + processPeersSequentially(discoverPeers, storeName, version, partition, tableFormat, resultFuture); return resultFuture; } @@ -121,6 +125,7 @@ public CompletionStage get(String storeName, int version, int parti * @param storeName the name of the store * @param version the version of the store * @param partition the partition of the store + * @param tableFormat the needed table format * @param resultFuture the future to complete with the InputStream of the blob */ private void processPeersSequentially( @@ -128,6 +133,7 @@ private void processPeersSequentially( String storeName, int version, int partition, + BlobTransferTableFormat tableFormat, CompletableFuture resultFuture) { String replicaId = Utils.getReplicaId(Version.composeKafkaTopic(storeName, version), partition); Instant startTime = Instant.now(); @@ -150,7 +156,7 @@ private void processPeersSequentially( // Attempt to fetch the blob from the current peer asynchronously LOGGER.info("Attempting to connect to host: {}", chosenHost); - return nettyClient.get(chosenHost, storeName, version, partition) + return nettyClient.get(chosenHost, storeName, version, partition, tableFormat) .toCompletableFuture() .thenAccept(inputStream -> { // Success case: Complete the future with the input stream diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/NettyFileTransferClient.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/NettyFileTransferClient.java index ebdced80985..d8584e53c6a 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/NettyFileTransferClient.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/NettyFileTransferClient.java @@ -1,5 +1,6 @@ package com.linkedin.davinci.blobtransfer.client; +import com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferTableFormat; import com.linkedin.davinci.storage.StorageMetadataService; import com.linkedin.venice.exceptions.VenicePeersConnectionException; import io.netty.bootstrap.Bootstrap; @@ -52,7 +53,12 @@ public void initChannel(SocketChannel ch) { }); } - public CompletionStage get(String host, String storeName, int version, int partition) { + public CompletionStage get( + String host, + String storeName, + int version, + int partition, + BlobTransferTableFormat requestedTableFormat) { CompletionStage inputStream = new CompletableFuture<>(); try { // Connects to the remote host @@ -64,10 +70,24 @@ public CompletionStage get(String host, String storeName, int versi ch.pipeline() .addLast(new IdleStateHandler(0, 0, 60)) .addLast(new MetadataAggregator(MAX_METADATA_CONTENT_LENGTH)) - .addLast(new P2PFileTransferClientHandler(baseDir, inputStream, storeName, version, partition)) - .addLast(new P2PMetadataTransferHandler(storageMetadataService, baseDir, storeName, version, partition)); + .addLast( + new P2PFileTransferClientHandler( + baseDir, + inputStream, + storeName, + version, + partition, + requestedTableFormat)) + .addLast( + new P2PMetadataTransferHandler( + storageMetadataService, + baseDir, + storeName, + version, + partition, + requestedTableFormat)); // Send a GET request - ch.writeAndFlush(prepareRequest(storeName, version, partition)); + ch.writeAndFlush(prepareRequest(storeName, version, partition, requestedTableFormat)); } catch (Exception e) { if (!inputStream.toCompletableFuture().isCompletedExceptionally()) { inputStream.toCompletableFuture().completeExceptionally(e); @@ -80,11 +100,15 @@ public void close() { workerGroup.shutdownGracefully(); } - private FullHttpRequest prepareRequest(String storeName, int version, int partition) { + private FullHttpRequest prepareRequest( + String storeName, + int version, + int partition, + BlobTransferTableFormat requestTableFormat) { return new DefaultFullHttpRequest( HttpVersion.HTTP_1_1, HttpMethod.GET, - String.format("/%s/%d/%d", storeName, version, partition)); + String.format("/%s/%d/%d/%s", storeName, version, partition, requestTableFormat.name())); } /** diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/P2PFileTransferClientHandler.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/P2PFileTransferClientHandler.java index accab2beea7..b87ea12db18 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/P2PFileTransferClientHandler.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/P2PFileTransferClientHandler.java @@ -56,9 +56,10 @@ public P2PFileTransferClientHandler( CompletionStage inputStreamFuture, String storeName, int version, - int partition) { + int partition, + BlobTransferUtils.BlobTransferTableFormat tableFormat) { this.inputStreamFuture = inputStreamFuture; - this.payload = new BlobTransferPayload(baseDir, storeName, version, partition); + this.payload = new BlobTransferPayload(baseDir, storeName, version, partition, tableFormat); } @Override diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/P2PMetadataTransferHandler.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/P2PMetadataTransferHandler.java index 32b6ca83186..3242b6d28a3 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/P2PMetadataTransferHandler.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/client/P2PMetadataTransferHandler.java @@ -4,6 +4,7 @@ import com.google.common.annotations.VisibleForTesting; import com.linkedin.davinci.blobtransfer.BlobTransferPartitionMetadata; import com.linkedin.davinci.blobtransfer.BlobTransferPayload; +import com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferTableFormat; import com.linkedin.davinci.storage.StorageMetadataService; import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.kafka.protocol.state.PartitionState; @@ -41,9 +42,10 @@ public P2PMetadataTransferHandler( String baseDir, String storeName, int version, - int partition) { + int partition, + BlobTransferTableFormat tableFormat) { this.storageMetadataService = storageMetadataService; - this.payload = new BlobTransferPayload(baseDir, storeName, version, partition); + this.payload = new BlobTransferPayload(baseDir, storeName, version, partition, tableFormat); } @Override diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/server/P2PFileTransferServerHandler.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/server/P2PFileTransferServerHandler.java index 903247a3b00..9cde393c77e 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/server/P2PFileTransferServerHandler.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/blobtransfer/server/P2PFileTransferServerHandler.java @@ -3,6 +3,7 @@ import static com.linkedin.davinci.blobtransfer.BlobTransferUtils.BLOB_TRANSFER_COMPLETED; import static com.linkedin.davinci.blobtransfer.BlobTransferUtils.BLOB_TRANSFER_STATUS; import static com.linkedin.davinci.blobtransfer.BlobTransferUtils.BLOB_TRANSFER_TYPE; +import static com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferTableFormat; import static com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferType; import static com.linkedin.venice.utils.NettyUtils.setupResponseAndFlush; import static io.netty.handler.codec.http.HttpHeaderValues.APPLICATION_JSON; @@ -121,6 +122,16 @@ protected void channelRead0(ChannelHandlerContext ctx, FullHttpRequest httpReque setupResponseAndFlush(HttpResponseStatus.NOT_FOUND, errBody, false, ctx); return; } + + // Check the snapshot table format + BlobTransferTableFormat currentSnapshotTableFormat = blobSnapshotManager.getBlobTransferTableFormat(); + if (blobTransferRequest.getRequestTableFormat() != currentSnapshotTableFormat) { + byte[] errBody = ("Table format mismatch for " + blobTransferRequest.getFullResourceName() + + ", current snapshot format is " + currentSnapshotTableFormat.name() + ", requested format is " + + blobTransferRequest.getRequestTableFormat().name()).getBytes(); + setupResponseAndFlush(HttpResponseStatus.NOT_FOUND, errBody, false, ctx); + return; + } } catch (IllegalArgumentException e) { setupResponseAndFlush(HttpResponseStatus.BAD_REQUEST, e.getMessage().getBytes(), false, ctx); return; @@ -278,13 +289,24 @@ public void sendMetadata(ChannelHandlerContext ctx, BlobTransferPartitionMetadat private BlobTransferPayload parseBlobTransferPayload(URI uri) throws IllegalArgumentException { // Parse the request uri to obtain the storeName and partition String[] requestParts = RequestHelper.getRequestParts(uri); - if (requestParts.length == 4) { - // [0]""/[1]"store"/[2]"version"/[3]"partition" + + // Ensure table format is valid + BlobTransferTableFormat requestTableFormat; + try { + requestTableFormat = BlobTransferTableFormat.valueOf(requestParts[4]); + } catch (IllegalArgumentException e) { + throw new IllegalArgumentException( + "Invalid table format: " + requestParts[4] + " for fetching blob at " + uri.getPath()); + } + + if (requestParts.length == 5) { + // [0]""/[1]"store"/[2]"version"/[3]"partition/[4]"table format" return new BlobTransferPayload( baseDir, requestParts[1], Integer.parseInt(requestParts[2]), - Integer.parseInt(requestParts[3])); + Integer.parseInt(requestParts[3]), + requestTableFormat); } else { throw new IllegalArgumentException("Invalid request for fetching blob at " + uri.getPath()); } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/AvroGenericDaVinciClient.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/AvroGenericDaVinciClient.java index 348b48128f2..a23c2dc8811 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/AvroGenericDaVinciClient.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/AvroGenericDaVinciClient.java @@ -732,7 +732,7 @@ protected void initBackend( Optional> managedClients, ICProvider icProvider, Optional cacheConfig, - DaVinciRecordTransformerFunctionalInterface recordTransformerFunction) { + DaVinciRecordTransformerConfig recordTransformerConfig) { synchronized (AvroGenericDaVinciClient.class) { if (daVinciBackend == null) { logger @@ -744,7 +744,7 @@ protected void initBackend( managedClients, icProvider, cacheConfig, - recordTransformerFunction), + recordTransformerConfig), backend -> { // Ensure that existing backend is fully closed before a new one can be created. synchronized (AvroGenericDaVinciClient.class) { @@ -782,13 +782,7 @@ public synchronized void start() { logger.info("Starting client, storeName=" + getStoreName()); VeniceConfigLoader configLoader = buildVeniceConfig(); Optional cacheConfig = Optional.ofNullable(daVinciConfig.getCacheConfig()); - initBackend( - clientConfig, - configLoader, - managedClients, - icProvider, - cacheConfig, - daVinciConfig.getRecordTransformerFunction()); + initBackend(clientConfig, configLoader, managedClients, icProvider, cacheConfig, recordTransformerConfig); try { getBackend().verifyCacheConfigEquality(daVinciConfig.getCacheConfig(), getStoreName()); diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/BlockingDaVinciRecordTransformer.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/BlockingDaVinciRecordTransformer.java index 9f4a2d92db2..85ae9f9168a 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/BlockingDaVinciRecordTransformer.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/BlockingDaVinciRecordTransformer.java @@ -19,19 +19,16 @@ public class BlockingDaVinciRecordTransformer extends DaVinciRecordTran private final DaVinciRecordTransformer recordTransformer; private final CountDownLatch startLatch = new CountDownLatch(1); - public BlockingDaVinciRecordTransformer(DaVinciRecordTransformer recordTransformer, boolean storeRecordsInDaVinci) { - super(recordTransformer.getStoreVersion(), storeRecordsInDaVinci); + public BlockingDaVinciRecordTransformer( + DaVinciRecordTransformer recordTransformer, + Schema keySchema, + Schema inputValueSchema, + Schema outputValueSchema, + boolean storeRecordsInDaVinci) { + super(recordTransformer.getStoreVersion(), keySchema, inputValueSchema, outputValueSchema, storeRecordsInDaVinci); this.recordTransformer = recordTransformer; } - public Schema getKeySchema() { - return this.recordTransformer.getKeySchema(); - } - - public Schema getOutputValueSchema() { - return this.recordTransformer.getOutputValueSchema(); - } - public DaVinciRecordTransformerResult transform(Lazy key, Lazy value) { return this.recordTransformer.transform(key, value); } @@ -51,12 +48,12 @@ public void processDelete(Lazy key) { this.recordTransformer.processDelete(key); } - public void onStartVersionIngestion() { - this.recordTransformer.onStartVersionIngestion(); + public void onStartVersionIngestion(boolean isCurrentVersion) { + this.recordTransformer.onStartVersionIngestion(isCurrentVersion); startLatch.countDown(); } - public void onEndVersionIngestion() { - this.recordTransformer.onEndVersionIngestion(); + public void onEndVersionIngestion(int currentVersion) { + this.recordTransformer.onEndVersionIngestion(currentVersion); } } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciConfig.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciConfig.java index a51e535e662..013c5f3c0db 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciConfig.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciConfig.java @@ -127,20 +127,6 @@ public DaVinciRecordTransformerConfig getRecordTransformerConfig() { return recordTransformerConfig; } - public DaVinciRecordTransformer getRecordTransformer(Integer storeVersion) { - if (recordTransformerConfig == null) { - return null; - } - return recordTransformerConfig.getRecordTransformer(storeVersion); - } - - public DaVinciRecordTransformerFunctionalInterface getRecordTransformerFunction() { - if (recordTransformerConfig == null) { - return null; - } - return recordTransformerConfig.getRecordTransformerFunction(); - } - public boolean isReadMetricsEnabled() { return readMetricsEnabled; } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformer.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformer.java index 526b42672dd..4a5a98c39e5 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformer.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformer.java @@ -40,33 +40,46 @@ public abstract class DaVinciRecordTransformer { */ private final boolean storeRecordsInDaVinci; + /** + * The key schema, which is immutable inside DaVinciClient. Users can modify the key if they are storing records in an external storage engine, but this must be managed by the user. + */ + private final Schema keySchema; + + /** + * The value schema before transformation, which is provided by the DaVinciClient. + */ + private final Schema inputValueSchema; + + /** + * The value schema after transformation, which is provided by the user. + */ + private final Schema outputValueSchema; + private final DaVinciRecordTransformerUtility recordTransformerUtility; /** * @param storeVersion the version of the store + * @param keySchema the key schema, which is immutable inside DaVinciClient. Users can modify the key if they are storing records in an external storage engine, but this must be managed by the user + * @param inputValueSchema the value schema before transformation + * @param outputValueSchema the value schema after transformation * @param storeRecordsInDaVinci set this to false if you intend to store records in a custom storage, - * and not in the Da Vinci Client. - */ - public DaVinciRecordTransformer(int storeVersion, boolean storeRecordsInDaVinci) { + * and not in the Da Vinci Client + */ + public DaVinciRecordTransformer( + int storeVersion, + Schema keySchema, + Schema inputValueSchema, + Schema outputValueSchema, + boolean storeRecordsInDaVinci) { this.storeVersion = storeVersion; this.storeRecordsInDaVinci = storeRecordsInDaVinci; + this.keySchema = keySchema; + // ToDo: Make use of inputValueSchema to support reader/writer schemas + this.inputValueSchema = inputValueSchema; + this.outputValueSchema = outputValueSchema; this.recordTransformerUtility = new DaVinciRecordTransformerUtility<>(this); } - /** - * Returns the schema for the key used in {@link DaVinciClient}'s operations. - * - * @return a {@link Schema} corresponding to the type of {@link K}. - */ - public abstract Schema getKeySchema(); - - /** - * Returns the schema for the output value used in {@link DaVinciClient}'s operations. - * - * @return a {@link Schema} corresponding to the type of {@link O}. - */ - public abstract Schema getOutputValueSchema(); - /** * Implement this method to transform records before they are stored. * This can be useful for tasks such as filtering out unused fields to save storage space. @@ -103,7 +116,7 @@ public void processDelete(Lazy key) { * * By default, it performs no operation. */ - public void onStartVersionIngestion() { + public void onStartVersionIngestion(boolean isCurrentVersion) { return; } @@ -113,10 +126,12 @@ public void onStartVersionIngestion() { * * By default, it performs no operation. */ - public void onEndVersionIngestion() { + public void onEndVersionIngestion(int currentVersion) { return; } + // Final methods below + /** * Transforms and processes the given record. * @@ -204,6 +219,33 @@ public final boolean getStoreRecordsInDaVinci() { return storeRecordsInDaVinci; } + /** + * Returns the schema for the key used in {@link DaVinciClient}'s operations. + * + * @return a {@link Schema} corresponding to the type of {@link K}. + */ + public final Schema getKeySchema() { + return keySchema; + } + + /** + * Returns the schema for the input value used in {@link DaVinciClient}'s operations. + * + * @return a {@link Schema} corresponding to the type of {@link V}. + */ + public final Schema getInputValueSchema() { + return inputValueSchema; + } + + /** + * Returns the schema for the output value used in {@link DaVinciClient}'s operations. + * + * @return a {@link Schema} corresponding to the type of {@link O}. + */ + public final Schema getOutputValueSchema() { + return outputValueSchema; + } + /** * @return {@link #recordTransformerUtility} */ diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerConfig.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerConfig.java index a92df1edae0..925aec2520c 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerConfig.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerConfig.java @@ -32,14 +32,6 @@ public DaVinciRecordTransformerFunctionalInterface getRecordTransformerFunction( return recordTransformerFunction; } - /** - * @param storeVersion the store version - * @return a new {@link DaVinciRecordTransformer} - */ - public DaVinciRecordTransformer getRecordTransformer(Integer storeVersion) { - return recordTransformerFunction.apply(storeVersion); - } - /** * @return {@link #outputValueClass} */ @@ -53,5 +45,4 @@ public Class getOutputValueClass() { public Schema getOutputValueSchema() { return outputValueSchema; } - } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerFunctionalInterface.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerFunctionalInterface.java index 9ad02fec641..e821d935ff2 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerFunctionalInterface.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerFunctionalInterface.java @@ -1,10 +1,17 @@ package com.linkedin.davinci.client; +import org.apache.avro.Schema; + + /** * This describes the implementation for the functional interface of {@link DaVinciRecordTransformer} */ @FunctionalInterface public interface DaVinciRecordTransformerFunctionalInterface { - DaVinciRecordTransformer apply(Integer storeVersion); + DaVinciRecordTransformer apply( + Integer storeVersion, + Schema keySchema, + Schema inputValueSchema, + Schema outputValueSchema); } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerUtility.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerUtility.java index 72cdf35a867..8732a633a5b 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerUtility.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/client/DaVinciRecordTransformerUtility.java @@ -24,12 +24,18 @@ */ public class DaVinciRecordTransformerUtility { private final DaVinciRecordTransformer recordTransformer; - private AvroGenericDeserializer keyDeserializer; - private AvroGenericDeserializer outputValueDeserializer; - private AvroSerializer outputValueSerializer; + private final AvroGenericDeserializer keyDeserializer; + private final AvroGenericDeserializer outputValueDeserializer; + private final AvroSerializer outputValueSerializer; public DaVinciRecordTransformerUtility(DaVinciRecordTransformer recordTransformer) { this.recordTransformer = recordTransformer; + + Schema keySchema = recordTransformer.getKeySchema(); + Schema outputValueSchema = recordTransformer.getOutputValueSchema(); + this.keyDeserializer = new AvroGenericDeserializer<>(keySchema, keySchema); + this.outputValueDeserializer = new AvroGenericDeserializer<>(outputValueSchema, outputValueSchema); + this.outputValueSerializer = new AvroSerializer<>(outputValueSchema); } /** @@ -40,7 +46,7 @@ public DaVinciRecordTransformerUtility(DaVinciRecordTransformer recordTransforme * @return a ByteBuffer containing the schema ID followed by the serialized and compressed value */ public final ByteBuffer prependSchemaIdToHeader(O value, int schemaId, VeniceCompressor compressor) { - byte[] serializedValue = getOutputValueSerializer().serialize(value); + byte[] serializedValue = outputValueSerializer.serialize(value); byte[] compressedValue; try { compressedValue = compressor.compress(serializedValue); @@ -116,7 +122,7 @@ public final void onRecovery( for (iterator.seekToFirst(); iterator.isValid(); iterator.next()) { byte[] keyBytes = iterator.key(); byte[] valueBytes = iterator.value(); - Lazy lazyKey = Lazy.of(() -> getKeyDeserializer().deserialize(ByteBuffer.wrap(keyBytes))); + Lazy lazyKey = Lazy.of(() -> keyDeserializer.deserialize(keyBytes)); Lazy lazyValue = Lazy.of(() -> { ByteBuffer valueByteBuffer = ByteBuffer.wrap(valueBytes); // Skip schema id @@ -127,35 +133,11 @@ public final void onRecovery( } catch (IOException e) { throw new RuntimeException(e); } - return getOutputValueDeserializer().deserialize(decompressedValueBytes); + return outputValueDeserializer.deserialize(decompressedValueBytes); }); recordTransformer.processPut(lazyKey, lazyValue); } } } - - public AvroGenericDeserializer getKeyDeserializer() { - if (keyDeserializer == null) { - Schema keySchema = recordTransformer.getKeySchema(); - keyDeserializer = new AvroGenericDeserializer<>(keySchema, keySchema); - } - return keyDeserializer; - } - - public AvroGenericDeserializer getOutputValueDeserializer() { - if (outputValueDeserializer == null) { - Schema outputValueSchema = recordTransformer.getOutputValueSchema(); - outputValueDeserializer = new AvroGenericDeserializer<>(outputValueSchema, outputValueSchema); - } - return outputValueDeserializer; - } - - public AvroSerializer getOutputValueSerializer() { - if (outputValueSerializer == null) { - Schema outputValueSchema = recordTransformer.getOutputValueSchema(); - outputValueSerializer = new AvroSerializer<>(outputValueSchema); - } - return outputValueSerializer; - } } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/config/VeniceServerConfig.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/config/VeniceServerConfig.java index bcc80a2cc9e..769b58c19ae 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/config/VeniceServerConfig.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/config/VeniceServerConfig.java @@ -56,6 +56,9 @@ import static com.linkedin.venice.ConfigKeys.SERVER_AA_WC_LEADER_QUOTA_RECORDS_PER_SECOND; import static com.linkedin.venice.ConfigKeys.SERVER_AA_WC_WORKLOAD_PARALLEL_PROCESSING_ENABLED; import static com.linkedin.venice.ConfigKeys.SERVER_AA_WC_WORKLOAD_PARALLEL_PROCESSING_THREAD_POOL_SIZE; +import static com.linkedin.venice.ConfigKeys.SERVER_ADAPTIVE_THROTTLER_ENABLED; +import static com.linkedin.venice.ConfigKeys.SERVER_ADAPTIVE_THROTTLER_SIGNAL_IDLE_THRESHOLD; +import static com.linkedin.venice.ConfigKeys.SERVER_ADAPTIVE_THROTTLER_SINGLE_GET_LATENCY_THRESHOLD; import static com.linkedin.venice.ConfigKeys.SERVER_BATCH_REPORT_END_OF_INCREMENTAL_PUSH_STATUS_ENABLED; import static com.linkedin.venice.ConfigKeys.SERVER_BLOCKING_QUEUE_TYPE; import static com.linkedin.venice.ConfigKeys.SERVER_CHANNEL_OPTION_WRITE_BUFFER_WATERMARK_HIGH_BYTES; @@ -496,6 +499,11 @@ public class VeniceServerConfig extends VeniceClusterConfig { private final boolean unregisterMetricForDeletedStoreEnabled; protected final boolean readOnlyForBatchOnlyStoreEnabled; // TODO: remove this config as its never used in prod private final boolean resetErrorReplicaEnabled; + + private final boolean adaptiveThrottlerEnabled; + private final int adaptiveThrottlerSignalIdleThreshold; + private final double adaptiveThrottlerSingleGetLatencyThreshold; + private final int fastAvroFieldLimitPerMethod; /** @@ -659,6 +667,11 @@ public VeniceServerConfig(VeniceProperties serverProperties, Map bootstrapFuture = bootstrapFromBlobs( storeAndVersion.getFirst(), storeAndVersion.getSecond().getNumber(), partition, + requestTableFormat, serverConfig.getBlobTransferDisabledOffsetLagThreshold()); bootstrapFuture.whenComplete((result, throwable) -> { @@ -103,6 +111,7 @@ CompletionStage bootstrapFromBlobs( Store store, int versionNumber, int partitionId, + BlobTransferTableFormat tableFormat, long blobTransferDisabledOffsetLagThreshold) { if (!store.isBlobTransferEnabled() || blobTransferManager == null) { return CompletableFuture.completedFuture(null); @@ -115,19 +124,23 @@ CompletionStage bootstrapFromBlobs( } String storeName = store.getName(); - return blobTransferManager.get(storeName, versionNumber, partitionId).handle((inputStream, throwable) -> { - updateBlobTransferResponseStats(throwable == null, storeName, versionNumber); - if (throwable != null) { - LOGGER.error( - "Failed to bootstrap partition {} from blobs transfer for store {} with exception {}", - partitionId, - storeName, - throwable); - } else { - LOGGER.info("Successfully bootstrapped partition {} from blobs transfer for store {}", partitionId, storeName); - } - return null; - }); + return blobTransferManager.get(storeName, versionNumber, partitionId, tableFormat) + .handle((inputStream, throwable) -> { + updateBlobTransferResponseStats(throwable == null, storeName, versionNumber); + if (throwable != null) { + LOGGER.error( + "Failed to bootstrap partition {} from blobs transfer for store {} with exception {}", + partitionId, + storeName, + throwable); + } else { + LOGGER.info( + "Successfully bootstrapped partition {} from blobs transfer for store {}", + partitionId, + storeName); + } + return null; + }); } /** diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/ingestion/isolated/IsolatedIngestionServer.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/ingestion/isolated/IsolatedIngestionServer.java index 925624a5b27..61a535c4455 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/ingestion/isolated/IsolatedIngestionServer.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/ingestion/isolated/IsolatedIngestionServer.java @@ -717,6 +717,7 @@ private void initializeIsolatedIngestionServer() { pubSubClientsFactory, sslFactory, null, + null, null); storeIngestionService.start(); storeIngestionService.addIngestionNotifier(new IsolatedIngestionNotifier(this)); diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/ActiveActiveStoreIngestionTask.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/ActiveActiveStoreIngestionTask.java index 51c59eadd6a..493b31bfc03 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/ActiveActiveStoreIngestionTask.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/ActiveActiveStoreIngestionTask.java @@ -5,7 +5,7 @@ import static com.linkedin.venice.writer.VeniceWriter.APP_DEFAULT_LOGICAL_TS; import com.linkedin.avroutil1.compatibility.AvroCompatibilityHelper; -import com.linkedin.davinci.client.DaVinciRecordTransformerFunctionalInterface; +import com.linkedin.davinci.client.DaVinciRecordTransformerConfig; import com.linkedin.davinci.config.VeniceServerConfig; import com.linkedin.davinci.config.VeniceStoreVersionConfig; import com.linkedin.davinci.replication.RmdWithValueSchemaId; @@ -114,7 +114,7 @@ public ActiveActiveStoreIngestionTask( int errorPartitionId, boolean isIsolatedIngestion, Optional cacheBackend, - DaVinciRecordTransformerFunctionalInterface recordTransformerFunction, + DaVinciRecordTransformerConfig recordTransformerConfig, Lazy zkHelixAdmin) { super( storageService, @@ -127,7 +127,7 @@ public ActiveActiveStoreIngestionTask( errorPartitionId, isIsolatedIngestion, cacheBackend, - recordTransformerFunction, + recordTransformerConfig, zkHelixAdmin); this.rmdProtocolVersionId = version.getRmdVersionId(); diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/AdaptiveThrottlerSignalService.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/AdaptiveThrottlerSignalService.java new file mode 100644 index 00000000000..674cc01d06f --- /dev/null +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/AdaptiveThrottlerSignalService.java @@ -0,0 +1,125 @@ +package com.linkedin.davinci.kafka.consumer; + +import com.linkedin.alpini.base.concurrency.Executors; +import com.linkedin.alpini.base.concurrency.ScheduledExecutorService; +import com.linkedin.davinci.config.VeniceServerConfig; +import com.linkedin.davinci.stats.ingestion.heartbeat.AggregatedHeartbeatLagEntry; +import com.linkedin.davinci.stats.ingestion.heartbeat.HeartbeatMonitoringService; +import com.linkedin.venice.service.AbstractVeniceService; +import io.tehuti.Metric; +import io.tehuti.metrics.MetricsRepository; +import java.util.ArrayList; +import java.util.List; +import java.util.concurrent.TimeUnit; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * This class contains service to periodically refresh all the signals for throttlers and update all registered throttler + * based on new signal values. + */ +public class AdaptiveThrottlerSignalService extends AbstractVeniceService { + public static final long HEARTBEAT_LAG_LIMIT = TimeUnit.MINUTES.toMillis(10); + public static final String SINGLE_GET_LATENCY_P99_METRIC_NAME = ".total--success_request_latency.99thPercentile"; + private static final Logger LOGGER = LogManager.getLogger(AdaptiveThrottlerSignalService.class); + private final double singleGetLatencyP99Threshold; + private final MetricsRepository metricsRepository; + private final HeartbeatMonitoringService heartbeatMonitoringService; + private final List throttlerList = new ArrayList<>(); + private final ScheduledExecutorService updateService = Executors.newSingleThreadScheduledExecutor(); + private boolean singleGetLatencySignal = false; + private boolean currentLeaderMaxHeartbeatLagSignal = false; + private boolean currentFollowerMaxHeartbeatLagSignal = false; + private boolean nonCurrentLeaderMaxHeartbeatLagSignal = false; + private boolean nonCurrentFollowerMaxHeartbeatLagSignal = false; + + public AdaptiveThrottlerSignalService( + VeniceServerConfig veniceServerConfig, + MetricsRepository metricsRepository, + HeartbeatMonitoringService heartbeatMonitoringService) { + this.singleGetLatencyP99Threshold = veniceServerConfig.getAdaptiveThrottlerSingleGetLatencyThreshold(); + this.metricsRepository = metricsRepository; + this.heartbeatMonitoringService = heartbeatMonitoringService; + } + + public void registerThrottler(VeniceAdaptiveIngestionThrottler adaptiveIngestionThrottler) { + throttlerList.add(adaptiveIngestionThrottler); + } + + public void refreshSignalAndThrottler() { + // Update all the signals in one shot; + updateReadLatencySignal(); + updateHeartbeatLatencySignal(); + // Update all the throttler + throttlerList.forEach(VeniceAdaptiveIngestionThrottler::checkSignalAndAdjustThrottler); + } + + void updateReadLatencySignal() { + Metric hostSingleGetLatencyP99Metric = metricsRepository.getMetric(SINGLE_GET_LATENCY_P99_METRIC_NAME); + if (hostSingleGetLatencyP99Metric != null) { + double hostSingleGetLatencyP99 = hostSingleGetLatencyP99Metric.value(); + singleGetLatencySignal = hostSingleGetLatencyP99 > singleGetLatencyP99Threshold; + LOGGER.info("Retrieved single get latency p99 value: {}", hostSingleGetLatencyP99); + } + LOGGER.info("Update read latency signal. singleGetLatency: {}", singleGetLatencySignal); + } + + void updateHeartbeatLatencySignal() { + AggregatedHeartbeatLagEntry maxLeaderHeartbeatLag = heartbeatMonitoringService.getMaxLeaderHeartbeatLag(); + if (maxLeaderHeartbeatLag != null) { + currentLeaderMaxHeartbeatLagSignal = maxLeaderHeartbeatLag.getCurrentVersionHeartbeatLag() > HEARTBEAT_LAG_LIMIT; + nonCurrentLeaderMaxHeartbeatLagSignal = + maxLeaderHeartbeatLag.getNonCurrentVersionHeartbeatLag() > HEARTBEAT_LAG_LIMIT; + } + AggregatedHeartbeatLagEntry maxFollowerHeartbeatLag = heartbeatMonitoringService.getMaxFollowerHeartbeatLag(); + if (maxFollowerHeartbeatLag != null) { + currentFollowerMaxHeartbeatLagSignal = + maxFollowerHeartbeatLag.getCurrentVersionHeartbeatLag() > HEARTBEAT_LAG_LIMIT; + nonCurrentFollowerMaxHeartbeatLagSignal = + maxFollowerHeartbeatLag.getNonCurrentVersionHeartbeatLag() > HEARTBEAT_LAG_LIMIT; + } + LOGGER.info( + "Update heartbeat signal. currentLeader: {}, currentFollower: {}, nonCurrentLeader: {}, nonCurrentFollower: {}", + currentLeaderMaxHeartbeatLagSignal, + currentFollowerMaxHeartbeatLagSignal, + nonCurrentLeaderMaxHeartbeatLagSignal, + nonCurrentFollowerMaxHeartbeatLagSignal); + + } + + public boolean isSingleGetLatencySignalActive() { + return singleGetLatencySignal; + } + + public boolean isCurrentLeaderMaxHeartbeatLagSignalActive() { + return currentLeaderMaxHeartbeatLagSignal; + } + + public boolean isCurrentFollowerMaxHeartbeatLagSignalActive() { + return currentFollowerMaxHeartbeatLagSignal; + } + + public boolean isNonCurrentLeaderMaxHeartbeatLagSignalActive() { + return nonCurrentLeaderMaxHeartbeatLagSignal; + } + + public boolean isNonCurrentFollowerMaxHeartbeatLagSignalActive() { + return nonCurrentFollowerMaxHeartbeatLagSignal; + } + + @Override + public boolean startInner() throws Exception { + updateService.scheduleAtFixedRate(this::refreshSignalAndThrottler, 1, 1, TimeUnit.MINUTES); + return true; + } + + @Override + public void stopInner() throws Exception { + updateService.shutdownNow(); + } + + List getThrottlerList() { + return throttlerList; + } +} diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/IngestionThrottler.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/IngestionThrottler.java index 2fe69cfab06..c648d0eeff3 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/IngestionThrottler.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/IngestionThrottler.java @@ -43,13 +43,15 @@ public class IngestionThrottler implements Closeable { public IngestionThrottler( boolean isDaVinciClient, VeniceServerConfig serverConfig, - Supplier> ongoingIngestionTaskMapSupplier) { + Supplier> ongoingIngestionTaskMapSupplier, + AdaptiveThrottlerSignalService adaptiveThrottlerSignalService) { this( isDaVinciClient, serverConfig, ongoingIngestionTaskMapSupplier, CURRENT_VERSION_BOOTSTRAPPING_DEFAULT_CHECK_INTERVAL, - CURRENT_VERSION_BOOTSTRAPPING_DEFAULT_CHECK_TIMEUNIT); + CURRENT_VERSION_BOOTSTRAPPING_DEFAULT_CHECK_TIMEUNIT, + adaptiveThrottlerSignalService); } public IngestionThrottler( @@ -57,15 +59,30 @@ public IngestionThrottler( VeniceServerConfig serverConfig, Supplier> ongoingIngestionTaskMapSupplier, int checkInterval, - TimeUnit checkTimeUnit) { - - EventThrottler regularRecordThrottler = new EventThrottler( - serverConfig.getKafkaFetchQuotaRecordPerSecond(), - serverConfig.getKafkaFetchQuotaTimeWindow(), - "kafka_consumption_records_count", - false, - EventThrottler.BLOCK_STRATEGY); - EventThrottler regularBandwidthThrottler = new EventThrottler( + TimeUnit checkTimeUnit, + AdaptiveThrottlerSignalService adaptiveThrottlerSignalService) { + VeniceAdaptiveIngestionThrottler globalRecordAdaptiveIngestionThrottler; + EventThrottler globalRecordThrottler; + if (serverConfig.isAdaptiveThrottlerEnabled()) { + globalRecordThrottler = null; + globalRecordAdaptiveIngestionThrottler = new VeniceAdaptiveIngestionThrottler( + serverConfig.getAdaptiveThrottlerSignalIdleThreshold(), + serverConfig.getKafkaFetchQuotaRecordPerSecond(), + serverConfig.getKafkaFetchQuotaTimeWindow(), + "kafka_consumption_records_count"); + globalRecordAdaptiveIngestionThrottler + .registerLimiterSignal(adaptiveThrottlerSignalService::isSingleGetLatencySignalActive); + adaptiveThrottlerSignalService.registerThrottler(globalRecordAdaptiveIngestionThrottler); + } else { + globalRecordAdaptiveIngestionThrottler = null; + globalRecordThrottler = new EventThrottler( + serverConfig.getKafkaFetchQuotaRecordPerSecond(), + serverConfig.getKafkaFetchQuotaTimeWindow(), + "kafka_consumption_records_count", + false, + EventThrottler.BLOCK_STRATEGY); + } + EventThrottler globalBandwidthThrottler = new EventThrottler( serverConfig.getKafkaFetchQuotaBytesPerSecond(), serverConfig.getKafkaFetchQuotaTimeWindow(), "kafka_consumption_bandwidth", @@ -157,8 +174,10 @@ public IngestionThrottler( this.isUsingSpeedupThrottler = true; } else if (!hasCurrentVersionBootstrapping && isUsingSpeedupThrottler) { LOGGER.info("There is no active current version bootstrapping, so switch to regular throttler"); - this.finalRecordThrottler = regularRecordThrottler; - this.finalBandwidthThrottler = regularBandwidthThrottler; + this.finalRecordThrottler = serverConfig.isAdaptiveThrottlerEnabled() + ? globalRecordAdaptiveIngestionThrottler + : globalRecordThrottler; + this.finalBandwidthThrottler = globalBandwidthThrottler; this.isUsingSpeedupThrottler = false; } @@ -170,8 +189,8 @@ public IngestionThrottler( this.eventThrottlerUpdateService = null; } - this.finalRecordThrottler = regularRecordThrottler; - this.finalBandwidthThrottler = regularBandwidthThrottler; + this.finalRecordThrottler = globalRecordThrottler; + this.finalBandwidthThrottler = globalBandwidthThrottler; } public void maybeThrottleRecordRate(ConsumerPoolType poolType, int count) { diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionService.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionService.java index ac12ed573a7..62384619ae0 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionService.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionService.java @@ -13,7 +13,7 @@ import static java.lang.Thread.currentThread; import static java.lang.Thread.sleep; -import com.linkedin.davinci.client.DaVinciRecordTransformerFunctionalInterface; +import com.linkedin.davinci.client.DaVinciRecordTransformerConfig; import com.linkedin.davinci.compression.StorageEngineBackedCompressorFactory; import com.linkedin.davinci.config.VeniceConfigLoader; import com.linkedin.davinci.config.VeniceServerConfig; @@ -177,7 +177,7 @@ public class KafkaStoreIngestionService extends AbstractVeniceService implements // source. This could be a view of the data, or in our case a cache, or both potentially. private final Optional cacheBackend; - private final DaVinciRecordTransformerFunctionalInterface recordTransformerFunction; + private final DaVinciRecordTransformerConfig recordTransformerConfig; private final PubSubProducerAdapterFactory producerAdapterFactory; @@ -191,6 +191,7 @@ public class KafkaStoreIngestionService extends AbstractVeniceService implements private KafkaValueSerializer kafkaValueSerializer; private final IngestionThrottler ingestionThrottler; private final ExecutorService aaWCWorkLoadProcessingThreadPool; + private final AdaptiveThrottlerSignalService adaptiveThrottlerSignalService; private VeniceServerConfig serverConfig; @@ -213,16 +214,17 @@ public KafkaStoreIngestionService( boolean isIsolatedIngestion, StorageEngineBackedCompressorFactory compressorFactory, Optional cacheBackend, - DaVinciRecordTransformerFunctionalInterface recordTransformerFunction, + DaVinciRecordTransformerConfig recordTransformerConfig, boolean isDaVinciClient, RemoteIngestionRepairService remoteIngestionRepairService, PubSubClientsFactory pubSubClientsFactory, Optional sslFactory, HeartbeatMonitoringService heartbeatMonitoringService, - Lazy zkHelixAdmin) { + Lazy zkHelixAdmin, + AdaptiveThrottlerSignalService adaptiveThrottlerSignalService) { this.storageService = storageService; this.cacheBackend = cacheBackend; - this.recordTransformerFunction = recordTransformerFunction; + this.recordTransformerConfig = recordTransformerConfig; this.storageMetadataService = storageMetadataService; this.metadataRepo = metadataRepo; this.topicNameToIngestionTaskMap = new ConcurrentSkipListMap<>(); @@ -243,11 +245,12 @@ public KafkaStoreIngestionService( new VeniceWriterFactory(veniceWriterProperties, producerAdapterFactory, metricsRepository); VeniceWriterFactory veniceWriterFactoryForMetaStoreWriter = new VeniceWriterFactory(veniceWriterProperties, producerAdapterFactory, null); - + this.adaptiveThrottlerSignalService = adaptiveThrottlerSignalService; this.ingestionThrottler = new IngestionThrottler( isDaVinciClient, serverConfig, - () -> Collections.unmodifiableMap(topicNameToIngestionTaskMap)); + () -> Collections.unmodifiableMap(topicNameToIngestionTaskMap), + adaptiveThrottlerSignalService); final Map kafkaUrlToRecordsThrottler; if (liveClusterConfigRepository != null) { @@ -535,7 +538,7 @@ private StoreIngestionTask createStoreIngestionTask( partitionId, isIsolatedIngestion, cacheBackend, - recordTransformerFunction, + recordTransformerConfig, zkHelixAdmin); } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/LeaderFollowerStoreIngestionTask.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/LeaderFollowerStoreIngestionTask.java index d327bb1fa4f..be4d833aca8 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/LeaderFollowerStoreIngestionTask.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/LeaderFollowerStoreIngestionTask.java @@ -15,7 +15,7 @@ import static java.lang.Long.max; import static java.util.concurrent.TimeUnit.MILLISECONDS; -import com.linkedin.davinci.client.DaVinciRecordTransformerFunctionalInterface; +import com.linkedin.davinci.client.DaVinciRecordTransformerConfig; import com.linkedin.davinci.config.VeniceStoreVersionConfig; import com.linkedin.davinci.helix.LeaderFollowerPartitionStateModel; import com.linkedin.davinci.ingestion.LagType; @@ -216,7 +216,7 @@ public LeaderFollowerStoreIngestionTask( int errorPartitionId, boolean isIsolatedIngestion, Optional cacheBackend, - DaVinciRecordTransformerFunctionalInterface recordTransformerFunction, + DaVinciRecordTransformerConfig recordTransformerConfig, Lazy zkHelixAdmin) { super( storageService, @@ -229,7 +229,7 @@ public LeaderFollowerStoreIngestionTask( errorPartitionId, isIsolatedIngestion, cacheBackend, - recordTransformerFunction, + recordTransformerConfig, builder.getLeaderFollowerNotifiers(), zkHelixAdmin); this.version = version; @@ -846,6 +846,7 @@ private boolean canSwitchToLeaderTopic(PartitionConsumptionState pcs) { private boolean isLocalVersionTopicPartitionFullyConsumed(PartitionConsumptionState pcs) { long localVTOff = pcs.getLatestProcessedLocalVersionTopicOffset(); long localVTEndOffset = getTopicPartitionEndOffSet(localKafkaServer, versionTopic, pcs.getPartition()); + if (localVTEndOffset == StatsErrorCode.LAG_MEASUREMENT_FAILURE.code) { return false; } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreBufferService.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreBufferService.java index 8f47975866c..fb0c1cabb3b 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreBufferService.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreBufferService.java @@ -65,6 +65,8 @@ public class StoreBufferService extends AbstractStoreBufferService { private final boolean isSorted; + private volatile boolean isStarted = false; + public StoreBufferService( int drainerNum, long bufferCapacityPerDrainer, @@ -308,12 +310,14 @@ public boolean startInner() { drainerList.add(drainer); } this.executorService.shutdown(); + isStarted = true; return true; } @Override public void stopInner() throws Exception { // Graceful shutdown + isStarted = false; drainerList.forEach(drainer -> drainer.stop()); if (this.executorService != null) { this.executorService.shutdownNow(); @@ -344,6 +348,10 @@ public long getMaxMemoryUsagePerDrainer() { long maxUsage = 0; boolean slowDrainerExists = false; + if (!isStarted) { + return maxUsage; + } + for (MemoryBoundBlockingQueue queue: blockingQueueArr) { maxUsage = Math.max(maxUsage, queue.getMemoryUsage()); if (queue.getMemoryUsage() > 0.8 * bufferCapacityPerDrainer) { diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java index da5a4d68709..64fab3eca4f 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTask.java @@ -12,6 +12,7 @@ import static com.linkedin.venice.ConfigKeys.KAFKA_BOOTSTRAP_SERVERS; import static com.linkedin.venice.LogMessages.KILLED_JOB_MESSAGE; import static com.linkedin.venice.kafka.protocol.enums.ControlMessageType.START_OF_SEGMENT; +import static com.linkedin.venice.pubsub.PubSubConstants.UNKNOWN_LATEST_OFFSET; import static com.linkedin.venice.utils.Utils.FATAL_DATA_VALIDATION_ERROR; import static com.linkedin.venice.utils.Utils.getReplicaId; import static java.util.Comparator.comparingInt; @@ -22,7 +23,7 @@ import com.linkedin.davinci.client.BlockingDaVinciRecordTransformer; import com.linkedin.davinci.client.DaVinciRecordTransformer; -import com.linkedin.davinci.client.DaVinciRecordTransformerFunctionalInterface; +import com.linkedin.davinci.client.DaVinciRecordTransformerConfig; import com.linkedin.davinci.client.DaVinciRecordTransformerResult; import com.linkedin.davinci.compression.StorageEngineBackedCompressorFactory; import com.linkedin.davinci.config.VeniceServerConfig; @@ -253,7 +254,7 @@ public abstract class StoreIngestionTask implements Runnable, Closeable { protected final Optional hybridStoreConfig; protected final Consumer divErrorMetricCallback; private final ExecutorService missingSOPCheckExecutor = Executors.newSingleThreadExecutor(); - private final VeniceStoreVersionConfig storeConfig; + private final VeniceStoreVersionConfig storeVersionConfig; protected final long readCycleDelayMs; protected final long emptyPollSleepMs; @@ -366,25 +367,27 @@ public StoreIngestionTask( Version version, Properties kafkaConsumerProperties, BooleanSupplier isCurrentVersion, - VeniceStoreVersionConfig storeConfig, + VeniceStoreVersionConfig storeVersionConfig, int errorPartitionId, boolean isIsolatedIngestion, Optional cacheBackend, - DaVinciRecordTransformerFunctionalInterface recordTransformerFunction, + DaVinciRecordTransformerConfig recordTransformerConfig, Queue notifiers, Lazy zkHelixAdmin) { - this.storeConfig = storeConfig; - this.readCycleDelayMs = storeConfig.getKafkaReadCycleDelayMs(); - this.emptyPollSleepMs = storeConfig.getKafkaEmptyPollSleepMs(); - this.databaseSyncBytesIntervalForTransactionalMode = storeConfig.getDatabaseSyncBytesIntervalForTransactionalMode(); - this.databaseSyncBytesIntervalForDeferredWriteMode = storeConfig.getDatabaseSyncBytesIntervalForDeferredWriteMode(); + this.storeVersionConfig = storeVersionConfig; + this.readCycleDelayMs = storeVersionConfig.getKafkaReadCycleDelayMs(); + this.emptyPollSleepMs = storeVersionConfig.getKafkaEmptyPollSleepMs(); + this.databaseSyncBytesIntervalForTransactionalMode = + storeVersionConfig.getDatabaseSyncBytesIntervalForTransactionalMode(); + this.databaseSyncBytesIntervalForDeferredWriteMode = + storeVersionConfig.getDatabaseSyncBytesIntervalForDeferredWriteMode(); this.kafkaProps = kafkaConsumerProperties; this.storageService = storageService; this.storageEngineRepository = builder.getStorageEngineRepository(); this.storageMetadataService = builder.getStorageMetadataService(); this.storeRepository = builder.getMetadataRepo(); this.schemaRepository = builder.getSchemaRepo(); - this.kafkaVersionTopic = storeConfig.getStoreVersionName(); + this.kafkaVersionTopic = storeVersionConfig.getStoreVersionName(); this.pubSubTopicRepository = builder.getPubSubTopicRepository(); this.versionTopic = pubSubTopicRepository.getTopic(kafkaVersionTopic); this.storeName = versionTopic.getStoreName(); @@ -414,13 +417,13 @@ public StoreIngestionTask( new KafkaDataIntegrityValidator(this.kafkaVersionTopic, DISABLED, producerStateMaxAgeMs); this.ingestionTaskName = String.format(CONSUMER_TASK_ID_FORMAT, kafkaVersionTopic); this.topicManagerRepository = builder.getTopicManagerRepository(); - this.readOnlyForBatchOnlyStoreEnabled = storeConfig.isReadOnlyForBatchOnlyStoreEnabled(); + this.readOnlyForBatchOnlyStoreEnabled = storeVersionConfig.isReadOnlyForBatchOnlyStoreEnabled(); this.hostLevelIngestionStats = builder.getIngestionStats().getStoreStats(storeName); this.versionedDIVStats = builder.getVersionedDIVStats(); this.versionedIngestionStats = builder.getVersionedStorageIngestionStats(); this.isRunning = new AtomicBoolean(true); this.emitMetrics = new AtomicBoolean(true); - this.resetErrorReplicaEnabled = storeConfig.isResetErrorReplicaEnabled(); + this.resetErrorReplicaEnabled = storeVersionConfig.isResetErrorReplicaEnabled(); this.storeBufferService = builder.getStoreBufferService(); this.isCurrentVersion = isCurrentVersion; @@ -470,11 +473,21 @@ public StoreIngestionTask( this.chunkAssembler = new ChunkAssembler(storeName); this.cacheBackend = cacheBackend; - if (recordTransformerFunction != null) { - DaVinciRecordTransformer clientRecordTransformer = recordTransformerFunction.apply(versionNumber); + if (recordTransformerConfig != null && recordTransformerConfig.getRecordTransformerFunction() != null) { + Schema keySchema = schemaRepository.getKeySchema(storeName).getSchema(); + Schema inputValueSchema = schemaRepository.getSupersetOrLatestValueSchema(storeName).getSchema(); + Schema outputValueSchema = recordTransformerConfig.getOutputValueSchema(); + + DaVinciRecordTransformer clientRecordTransformer = recordTransformerConfig.getRecordTransformerFunction() + .apply(versionNumber, keySchema, inputValueSchema, outputValueSchema); + this.recordTransformer = new BlockingDaVinciRecordTransformer( clientRecordTransformer, + keySchema, + inputValueSchema, + outputValueSchema, clientRecordTransformer.getStoreRecordsInDaVinci()); + versionedIngestionStats.registerTransformerLatencySensor(storeName, versionNumber); versionedIngestionStats.registerTransformerLifecycleStartLatency(storeName, versionNumber); versionedIngestionStats.registerTransformerLifecycleEndLatency(storeName, versionNumber); @@ -483,7 +496,7 @@ public StoreIngestionTask( // onStartVersionIngestion called here instead of run() because this needs to finish running // before bootstrapping starts long startTime = System.currentTimeMillis(); - recordTransformer.onStartVersionIngestion(); + recordTransformer.onStartVersionIngestion(isCurrentVersion.getAsBoolean()); long endTime = System.currentTimeMillis(); versionedIngestionStats.recordTransformerLifecycleStartLatency( storeName, @@ -543,7 +556,7 @@ public StoreIngestionTask( } this.batchReportIncPushStatusEnabled = !isDaVinciClient && serverConfig.getBatchReportEOIPEnabled(); this.parallelProcessingThreadPool = builder.getAAWCWorkLoadProcessingThreadPool(); - this.hostName = Utils.getHostName() + "_" + storeConfig.getListenerPort(); + this.hostName = Utils.getHostName() + "_" + storeVersionConfig.getListenerPort(); this.zkHelixAdmin = zkHelixAdmin; } @@ -689,7 +702,7 @@ public CompletableFuture dropStoragePartitionGracefully(PubSubTopicPartiti private void dropPartitionSynchronously(PubSubTopicPartition topicPartition) { LOGGER.info("{} Dropping partition: {}", ingestionTaskName, topicPartition); int partition = topicPartition.getPartitionNumber(); - this.storageService.dropStorePartition(storeConfig, partition, true); + this.storageService.dropStorePartition(storeVersionConfig, partition, true); LOGGER.info("{} Dropped partition: {}", ingestionTaskName, topicPartition); } @@ -2344,16 +2357,17 @@ private void reportStoreVersionTopicOffsetRewindMetrics(PartitionConsumptionStat * written to, the end offset is 0. */ protected long getTopicPartitionEndOffSet(String kafkaUrl, PubSubTopic pubSubTopic, int partition) { - long offsetFromConsumer = aggKafkaConsumerService - .getLatestOffsetBasedOnMetrics(kafkaUrl, versionTopic, new PubSubTopicPartitionImpl(pubSubTopic, partition)); + PubSubTopicPartition topicPartition = new PubSubTopicPartitionImpl(pubSubTopic, partition); + long offsetFromConsumer = + aggKafkaConsumerService.getLatestOffsetBasedOnMetrics(kafkaUrl, versionTopic, topicPartition); if (offsetFromConsumer >= 0) { return offsetFromConsumer; } try { return RetryUtils.executeWithMaxAttemptAndExponentialBackoff(() -> { long offset = getTopicManager(kafkaUrl).getLatestOffsetCachedNonBlocking(pubSubTopic, partition); - if (offset == -1) { - throw new VeniceException("Found latest offset -1"); + if (offset == UNKNOWN_LATEST_OFFSET) { + throw new VeniceException("Latest offset is unknown. Check if the topic: " + topicPartition + " exists."); } return offset; }, @@ -2894,7 +2908,33 @@ private void processStartOfPush( /* * Notify the underlying store engine about starting batch push. */ - beginBatchWrite(partition, startOfPush.sorted, partitionConsumptionState); + final boolean sorted; + if (serverConfig.getRocksDBServerConfig().isBlobFilesEnabled() && isHybridMode()) { + /** + * We would like to skip {@link com.linkedin.davinci.store.rocksdb.RocksDBSstFileWriter} for hybrid stores + * when RocksDB blob mode is enabled and here are the reasons: + * 1. Hybrid stores will use the same amount of MemTables eventually even not in batch processing phase. + * 2. SSTFileWriter + RocksDB blob mode will introduce additional space overhead in the following way: + * a. SSTFileWriter doesn't support RocksDB blob mode, which means even with blob enabled, SSTFileWriter + * will still write both key and value into the same SST file regardless of value size. + * b. When RocksDB ingests the generated SST files, it will put them in the bottom level. + * c. After finishing the batch portion, RocksDB won't use SSTFileWriter anymore and in the regular mode, when + * RocksDB blob mode is enabled, RocksDB will store the value in blob files when the value size is + * larger than the configured threshold, and this also means the LSM tree built by the real-time writes + * will be much smaller as it contains keys and smaller values/value pointers. + * d. As LSM tree is small, it is not easy to trigger a compaction in the bottom level (the bottom - 1 level + * needs to keep enough data to trigger the bottom-level compaction), so the staled entries in the bottom + * level will remain for a long time. + * e. RocksDB blob config tuning won't affect the large bottom-level SST files. + * 3. If we disable SSTFileWriter for hybrid stores when RocksDB blob mode is enabled, all the writes will + * go through MemTable and key/value separation logic will apply all the time, and blob related configs + * will apply to all the writes. + */ + sorted = false; + } else { + sorted = startOfPush.sorted; + } + beginBatchWrite(partition, sorted, partitionConsumptionState); partitionConsumptionState.setStartOfPushTimestamp(startOfPushKME.producerMetadata.messageTimestamp); ingestionNotificationDispatcher.reportStarted(partitionConsumptionState); @@ -2902,7 +2942,7 @@ private void processStartOfPush( if (previousStoreVersionState == null) { // No other partition of the same topic has started yet, let's initialize the StoreVersionState StoreVersionState newStoreVersionState = new StoreVersionState(); - newStoreVersionState.sorted = startOfPush.sorted; + newStoreVersionState.sorted = sorted; newStoreVersionState.chunked = startOfPush.chunked; newStoreVersionState.compressionStrategy = startOfPush.compressionStrategy; newStoreVersionState.compressionDictionary = startOfPush.compressionDictionary; @@ -2920,7 +2960,7 @@ private void processStartOfPush( StoreVersionState.class.getSimpleName(), kafkaVersionTopic); return newStoreVersionState; - } else if (previousStoreVersionState.sorted != startOfPush.sorted) { + } else if (previousStoreVersionState.sorted != sorted) { // Something very wrong is going on ): ... throw new VeniceException( "Unexpected: received multiple " + ControlMessageType.START_OF_PUSH.name() @@ -2979,7 +3019,7 @@ protected void processEndOfPush( /** * Generate snapshot after batch write is done. */ - if (storeConfig.isBlobTransferEnabled()) { + if (storeVersionConfig.isBlobTransferEnabled() && serverConfig.isBlobTransferManagerEnabled()) { storageEngine.createSnapshot(storagePartitionConfig); } @@ -3101,6 +3141,13 @@ private boolean processControlMessage( processEndOfIncrementalPush(controlMessage, partitionConsumptionState); break; case TOPIC_SWITCH: + TopicSwitch topicSwitch = (TopicSwitch) controlMessage.controlMessageUnion; + LOGGER.info( + "Received {} control message. Replica: {}, Offset: {} NewSource: {}", + type.name(), + partitionConsumptionState.getReplicaId(), + offset, + topicSwitch.getSourceKafkaServers()); checkReadyToServeAfterProcess = processTopicSwitch(controlMessage, partition, offset, partitionConsumptionState); break; @@ -4037,7 +4084,8 @@ public synchronized void close() { if (recordTransformer != null) { long startTime = System.currentTimeMillis(); - recordTransformer.onEndVersionIngestion(); + Store store = storeRepository.getStoreOrThrow(storeName); + recordTransformer.onEndVersionIngestion(store.getCurrentVersion()); long endTime = System.currentTimeMillis(); versionedIngestionStats.recordTransformerLifecycleEndLatency( storeName, diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTaskFactory.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTaskFactory.java index 93211b6f24e..81514bf10f7 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTaskFactory.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTaskFactory.java @@ -1,6 +1,6 @@ package com.linkedin.davinci.kafka.consumer; -import com.linkedin.davinci.client.DaVinciRecordTransformerFunctionalInterface; +import com.linkedin.davinci.client.DaVinciRecordTransformerConfig; import com.linkedin.davinci.compression.StorageEngineBackedCompressorFactory; import com.linkedin.davinci.config.VeniceServerConfig; import com.linkedin.davinci.config.VeniceStoreVersionConfig; @@ -55,7 +55,7 @@ public StoreIngestionTask getNewIngestionTask( int partitionId, boolean isIsolatedIngestion, Optional cacheBackend, - DaVinciRecordTransformerFunctionalInterface recordTransformerFunction, + DaVinciRecordTransformerConfig recordTransformerConfig, Lazy zkHelixAdmin) { if (version.isActiveActiveReplicationEnabled()) { return new ActiveActiveStoreIngestionTask( @@ -69,7 +69,7 @@ public StoreIngestionTask getNewIngestionTask( partitionId, isIsolatedIngestion, cacheBackend, - recordTransformerFunction, + recordTransformerConfig, zkHelixAdmin); } return new LeaderFollowerStoreIngestionTask( @@ -83,7 +83,7 @@ public StoreIngestionTask getNewIngestionTask( partitionId, isIsolatedIngestion, cacheBackend, - recordTransformerFunction, + recordTransformerConfig, zkHelixAdmin); } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/VeniceAdaptiveIngestionThrottler.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/VeniceAdaptiveIngestionThrottler.java new file mode 100644 index 00000000000..757fedc6159 --- /dev/null +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/kafka/consumer/VeniceAdaptiveIngestionThrottler.java @@ -0,0 +1,111 @@ +package com.linkedin.davinci.kafka.consumer; + +import com.linkedin.venice.throttle.EventThrottler; +import java.text.DecimalFormat; +import java.util.ArrayList; +import java.util.List; +import java.util.function.BooleanSupplier; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +public class VeniceAdaptiveIngestionThrottler extends EventThrottler { + private static final Logger LOGGER = LogManager.getLogger(VeniceAdaptiveIngestionThrottler.class); + private final List limiterSuppliers = new ArrayList<>(); + private final List boosterSuppliers = new ArrayList<>(); + + private final static int MAX_THROTTLERS = 7; + private final List eventThrottlers = new ArrayList<>(MAX_THROTTLERS); + private final int signalIdleThreshold; + private int signalIdleCount = 0; + private int currentThrottlerIndex = MAX_THROTTLERS / 2; + + public VeniceAdaptiveIngestionThrottler( + int signalIdleThreshold, + long quotaPerSecond, + long timeWindow, + String throttlerName) { + this.signalIdleThreshold = signalIdleThreshold; + double factor = 0.4; + DecimalFormat decimalFormat = new DecimalFormat("0.0"); + for (int i = 0; i < MAX_THROTTLERS; i++) { + EventThrottler eventThrottler = new EventThrottler( + (long) (quotaPerSecond * factor), + timeWindow, + throttlerName + decimalFormat.format(factor), + false, + EventThrottler.BLOCK_STRATEGY); + eventThrottlers.add(eventThrottler); + factor = factor + 0.2; + } + } + + @Override + public void maybeThrottle(double eventsSeen) { + eventThrottlers.get(currentThrottlerIndex).maybeThrottle(eventsSeen); + } + + public void registerLimiterSignal(BooleanSupplier supplier) { + limiterSuppliers.add(supplier); + } + + public void registerBoosterSignal(BooleanSupplier supplier) { + boosterSuppliers.add(supplier); + } + + public void checkSignalAndAdjustThrottler() { + boolean isSignalIdle = true; + boolean hasLimitedRate = false; + for (BooleanSupplier supplier: limiterSuppliers) { + if (supplier.getAsBoolean()) { + hasLimitedRate = true; + isSignalIdle = false; + signalIdleCount = 0; + if (currentThrottlerIndex > 0) { + currentThrottlerIndex--; + } + LOGGER.info( + "Found active limiter signal, adjusting throttler index to: {} with throttle rate: {}", + currentThrottlerIndex, + eventThrottlers.get(currentThrottlerIndex).getMaxRatePerSecond()); + } + } + // If any limiter signal is true do not booster the throttler + if (hasLimitedRate) { + return; + } + + for (BooleanSupplier supplier: boosterSuppliers) { + if (supplier.getAsBoolean()) { + isSignalIdle = false; + signalIdleCount = 0; + if (currentThrottlerIndex < MAX_THROTTLERS - 1) { + currentThrottlerIndex++; + } + LOGGER.info( + "Found active booster signal, adjusting throttler index to: {} with throttle rate: {}", + currentThrottlerIndex, + eventThrottlers.get(currentThrottlerIndex).getMaxRatePerSecond()); + } + } + if (isSignalIdle) { + signalIdleCount += 1; + LOGGER.info("No active signal found, increasing idle count to {}/{}", signalIdleCount, signalIdleThreshold); + if (signalIdleCount == signalIdleThreshold) { + if (currentThrottlerIndex < MAX_THROTTLERS - 1) { + currentThrottlerIndex++; + } + LOGGER.info( + "Reach max signal idle count, adjusting throttler index to: {} with throttle rate: {}", + currentThrottlerIndex, + eventThrottlers.get(currentThrottlerIndex).getMaxRatePerSecond()); + signalIdleCount = 0; + } + } + } + + // TEST + int getCurrentThrottlerIndex() { + return currentThrottlerIndex; + } +} diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/NativeMetadataRepository.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/NativeMetadataRepository.java index bf7fffe070a..2a215ab4afc 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/NativeMetadataRepository.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/NativeMetadataRepository.java @@ -1,8 +1,6 @@ package com.linkedin.davinci.repository; import static com.linkedin.venice.ConfigKeys.CLIENT_SYSTEM_STORE_REPOSITORY_REFRESH_INTERVAL_SECONDS; -import static com.linkedin.venice.system.store.MetaStoreWriter.KEY_STRING_SCHEMA_ID; -import static com.linkedin.venice.system.store.MetaStoreWriter.KEY_STRING_STORE_NAME; import static java.lang.Thread.currentThread; import com.linkedin.davinci.stats.NativeMetadataRepositoryStats; @@ -25,17 +23,11 @@ import com.linkedin.venice.schema.rmd.RmdSchemaEntry; import com.linkedin.venice.schema.writecompute.DerivedSchemaEntry; import com.linkedin.venice.service.ICProvider; -import com.linkedin.venice.system.store.MetaStoreDataType; -import com.linkedin.venice.systemstore.schemas.StoreClusterConfig; -import com.linkedin.venice.systemstore.schemas.StoreMetaKey; -import com.linkedin.venice.systemstore.schemas.StoreMetaValue; import com.linkedin.venice.utils.VeniceProperties; import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap; import java.time.Clock; import java.util.ArrayList; import java.util.Collection; -import java.util.Collections; -import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.Set; @@ -70,7 +62,7 @@ public abstract class NativeMetadataRepository private final Map storeConfigMap = new VeniceConcurrentHashMap<>(); // Local cache for key/value schemas. SchemaData supports one key schema per store only, which may need to be changed // for key schema evolvability. - private final Map schemaMap = new VeniceConcurrentHashMap<>(); + protected final Map schemaMap = new VeniceConcurrentHashMap<>(); private final ScheduledExecutorService scheduler = Executors.newScheduledThreadPool(1); private final Set listeners = new CopyOnWriteArraySet<>(); private final AtomicLong totalStoreReadQuota = new AtomicLong(); @@ -181,7 +173,7 @@ public Store refreshOneStore(String storeName) { } Store newStore = fetchStoreFromRemote(storeName, storeConfig.getCluster()); putStore(newStore); - getAndCacheSchemaDataFromSystemStore(storeName); + getAndCacheSchemaData(storeName); nativeMetadataRepositoryStats.updateCacheTimestamp(storeName, clock.millis()); return newStore; } catch (ServiceDiscoveryException | MissingKeyInStoreMetadataException e) { @@ -401,66 +393,7 @@ protected StoreConfig cacheStoreConfigFromRemote(String storeName) { protected abstract Store fetchStoreFromRemote(String storeName, String clusterName); - protected abstract StoreMetaValue getStoreMetaValue(String storeName, StoreMetaKey key); - - // Helper function with common code for retrieving StoreConfig from meta system store. - protected StoreConfig getStoreConfigFromMetaSystemStore(String storeName) { - StoreClusterConfig clusterConfig = getStoreMetaValue( - storeName, - MetaStoreDataType.STORE_CLUSTER_CONFIG - .getStoreMetaKey(Collections.singletonMap(KEY_STRING_STORE_NAME, storeName))).storeClusterConfig; - return new StoreConfig(clusterConfig); - } - - // Helper function with common code for retrieving SchemaData from meta system store. - protected SchemaData getSchemaData(String storeName) { - SchemaData schemaData = schemaMap.get(storeName); - SchemaEntry keySchema; - if (schemaData == null) { - // Retrieve the key schema and initialize SchemaData only if it's not cached yet. - StoreMetaKey keySchemaKey = MetaStoreDataType.STORE_KEY_SCHEMAS - .getStoreMetaKey(Collections.singletonMap(KEY_STRING_STORE_NAME, storeName)); - Map keySchemaMap = - getStoreMetaValue(storeName, keySchemaKey).storeKeySchemas.keySchemaMap; - if (keySchemaMap.isEmpty()) { - throw new VeniceException("No key schema found for store: " + storeName); - } - Map.Entry keySchemaEntry = keySchemaMap.entrySet().iterator().next(); - keySchema = - new SchemaEntry(Integer.parseInt(keySchemaEntry.getKey().toString()), keySchemaEntry.getValue().toString()); - schemaData = new SchemaData(storeName, keySchema); - } - StoreMetaKey valueSchemaKey = MetaStoreDataType.STORE_VALUE_SCHEMAS - .getStoreMetaKey(Collections.singletonMap(KEY_STRING_STORE_NAME, storeName)); - Map valueSchemaMap = - getStoreMetaValue(storeName, valueSchemaKey).storeValueSchemas.valueSchemaMap; - // Check the value schema string, if it's empty then try to query the other key space for individual value schema. - for (Map.Entry entry: valueSchemaMap.entrySet()) { - // Check if we already have the corresponding value schema - int valueSchemaId = Integer.parseInt(entry.getKey().toString()); - if (schemaData.getValueSchema(valueSchemaId) != null) { - continue; - } - if (entry.getValue().toString().isEmpty()) { - // The value schemas might be too large to be stored in a single K/V. - StoreMetaKey individualValueSchemaKey = - MetaStoreDataType.STORE_VALUE_SCHEMA.getStoreMetaKey(new HashMap() { - { - put(KEY_STRING_STORE_NAME, storeName); - put(KEY_STRING_SCHEMA_ID, entry.getKey().toString()); - } - }); - // Empty string is not a valid value schema therefore it's safe to throw exceptions if we also cannot find it in - // the individual value schema key space. - String valueSchema = - getStoreMetaValue(storeName, individualValueSchemaKey).storeValueSchema.valueSchema.toString(); - schemaData.addValueSchema(new SchemaEntry(valueSchemaId, valueSchema)); - } else { - schemaData.addValueSchema(new SchemaEntry(valueSchemaId, entry.getValue().toString())); - } - } - return schemaData; - } + protected abstract SchemaData getSchemaData(String storeName); protected Store putStore(Store newStore) { // Workaround to make old metadata compatible with new fields @@ -516,7 +449,7 @@ protected void notifyStoreChanged(Store store) { } } - protected SchemaData getAndCacheSchemaDataFromSystemStore(String storeName) { + protected SchemaData getAndCacheSchemaData(String storeName) { if (!hasStore(storeName)) { throw new VeniceNoStoreException(storeName); } @@ -532,7 +465,7 @@ protected SchemaData getAndCacheSchemaDataFromSystemStore(String storeName) { private SchemaData getSchemaDataFromReadThroughCache(String storeName) throws VeniceNoStoreException { SchemaData schemaData = schemaMap.get(storeName); if (schemaData == null) { - schemaData = getAndCacheSchemaDataFromSystemStore(storeName); + schemaData = getAndCacheSchemaData(storeName); } return schemaData; } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/RequestBasedMetaRepository.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/RequestBasedMetaRepository.java index 4c9017ba97c..7a276f81e8c 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/RequestBasedMetaRepository.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/RequestBasedMetaRepository.java @@ -9,11 +9,11 @@ import com.linkedin.venice.meta.StoreConfig; import com.linkedin.venice.meta.ZKStore; import com.linkedin.venice.metadata.response.StorePropertiesResponseRecord; +import com.linkedin.venice.schema.SchemaData; +import com.linkedin.venice.schema.SchemaEntry; import com.linkedin.venice.serializer.FastSerializerDeserializerFactory; import com.linkedin.venice.serializer.RecordDeserializer; import com.linkedin.venice.systemstore.schemas.StoreClusterConfig; -import com.linkedin.venice.systemstore.schemas.StoreMetaKey; -import com.linkedin.venice.systemstore.schemas.StoreMetaValue; import com.linkedin.venice.systemstore.schemas.StoreProperties; import com.linkedin.venice.utils.VeniceProperties; import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap; @@ -22,11 +22,13 @@ public class RequestBasedMetaRepository extends NativeMetadataRepository { - // storeName -> storePropertiesResponseRecord - private final Map storePropertiesRecordMap = new VeniceConcurrentHashMap<>(); - private final Map d2TransportClientMap = new VeniceConcurrentHashMap<>(); // cluster -> - // client + // cluster -> client + private final Map d2TransportClientMap = new VeniceConcurrentHashMap<>(); + + // storeName -> T + private final Map storeSchemaMap = new VeniceConcurrentHashMap<>(); + private final D2TransportClient d2DiscoveryTransportClient; private D2ServiceDiscovery d2ServiceDiscovery; @@ -40,7 +42,10 @@ public RequestBasedMetaRepository(ClientConfig clientConfig, VeniceProperties ba @Override public void clear() { super.clear(); - storePropertiesRecordMap.clear(); + + // Clear cache + d2TransportClientMap.clear(); + storeSchemaMap.clear(); } @Override @@ -64,31 +69,32 @@ protected Store fetchStoreFromRemote(String storeName, String clusterName) { } @Override - protected StoreMetaValue getStoreMetaValue(String storeName, StoreMetaKey key) { // TODO PRANAV what is key for? - System.out.println("HERE HERE HERE getStoreMetaValue - Key: " + key); - if (storePropertiesRecordMap.containsKey(storeName)) { - return storePropertiesRecordMap.get(storeName).storeMetaValue; + protected SchemaData getSchemaData(String storeName) { + if (!storeSchemaMap.containsKey(storeName)) { + // Cache miss + fetchAndCacheStorePropertiesResponseRecord(storeName); } - // Cache miss, fetch store - StorePropertiesResponseRecord record = fetchAndCacheStorePropertiesResponseRecord(storeName); - return record.storeMetaValue; + return storeSchemaMap.get(storeName); } private StorePropertiesResponseRecord fetchAndCacheStorePropertiesResponseRecord(String storeName) { // Request - // TODO PRANAV lastKnownSchemaId param + int maxValueSchemaId = getMaxValueSchemaId(storeName); D2TransportClient d2TransportClient = getD2TransportClient(storeName); String requestBasedStorePropertiesURL = QueryAction.STORE_PROPERTIES.toString().toLowerCase() + "/" + storeName; + if (maxValueSchemaId > SchemaData.UNKNOWN_SCHEMA_ID) { + requestBasedStorePropertiesURL += "/" + maxValueSchemaId; + } + TransportClientResponse response; try { response = d2TransportClient.get(requestBasedStorePropertiesURL).get(); } catch (Exception e) { throw new RuntimeException( - "Encountered exception while trying to send store properties request to: " + requestBasedStorePropertiesURL - + "/" + requestBasedStorePropertiesURL, - e); + "Encountered exception while trying to send store properties request to " + requestBasedStorePropertiesURL + + ": " + e); } // Deserialize @@ -98,7 +104,7 @@ private StorePropertiesResponseRecord fetchAndCacheStorePropertiesResponseRecord StorePropertiesResponseRecord record = recordDeserializer.deserialize(response.getBody()); // Cache - storePropertiesRecordMap.put(storeName, record); + cacheStoreSchema(storeName, record); return record; } @@ -116,4 +122,33 @@ D2TransportClient getD2TransportClient(String storeName) { return d2TransportClient; } } + + private int getMaxValueSchemaId(String storeName) { + if (!schemaMap.containsKey(storeName)) { + return SchemaData.UNKNOWN_SCHEMA_ID; + } + return schemaMap.get(storeName).getMaxValueSchemaId(); + } + + private void cacheStoreSchema(String storeName, StorePropertiesResponseRecord record) { + + if (!storeSchemaMap.containsKey(storeName)) { + // New schema data + Map.Entry keySchemaEntry = + record.getStoreMetaValue().getStoreKeySchemas().getKeySchemaMap().entrySet().iterator().next(); + SchemaData schemaData = new SchemaData( + storeName, + new SchemaEntry(Integer.parseInt(keySchemaEntry.getKey().toString()), keySchemaEntry.getValue().toString())); + storeSchemaMap.put(storeName, schemaData); + } + + // Store Value Schemas + for (Map.Entry entry: record.getStoreMetaValue() + .getStoreValueSchemas() + .getValueSchemaMap() + .entrySet()) { + storeSchemaMap.get(storeName) + .addValueSchema(new SchemaEntry(Integer.parseInt(entry.getKey().toString()), entry.getValue().toString())); + } + } } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/ThinClientMetaStoreBasedRepository.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/ThinClientMetaStoreBasedRepository.java index 29758508406..b1caf645ccd 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/ThinClientMetaStoreBasedRepository.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/repository/ThinClientMetaStoreBasedRepository.java @@ -1,7 +1,6 @@ package com.linkedin.davinci.repository; -import static com.linkedin.venice.system.store.MetaStoreWriter.KEY_STRING_CLUSTER_NAME; -import static com.linkedin.venice.system.store.MetaStoreWriter.KEY_STRING_STORE_NAME; +import static com.linkedin.venice.system.store.MetaStoreWriter.*; import com.linkedin.venice.client.exceptions.ServiceDiscoveryException; import com.linkedin.venice.client.store.AvroSpecificStoreClient; @@ -9,12 +8,16 @@ import com.linkedin.venice.client.store.ClientFactory; import com.linkedin.venice.common.VeniceSystemStoreType; import com.linkedin.venice.exceptions.MissingKeyInStoreMetadataException; +import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.exceptions.VeniceRetriableException; import com.linkedin.venice.meta.Store; import com.linkedin.venice.meta.StoreConfig; import com.linkedin.venice.meta.ZKStore; +import com.linkedin.venice.schema.SchemaData; +import com.linkedin.venice.schema.SchemaEntry; import com.linkedin.venice.service.ICProvider; import com.linkedin.venice.system.store.MetaStoreDataType; +import com.linkedin.venice.systemstore.schemas.StoreClusterConfig; import com.linkedin.venice.systemstore.schemas.StoreMetaKey; import com.linkedin.venice.systemstore.schemas.StoreMetaValue; import com.linkedin.venice.systemstore.schemas.StoreProperties; @@ -78,7 +81,6 @@ protected Store fetchStoreFromRemote(String storeName, String clusterName) { return new ZKStore(storeProperties); } - @Override protected StoreMetaValue getStoreMetaValue(String storeName, StoreMetaKey key) { final Callable> supplier = () -> getAvroClientForMetaStore(storeName).get(key); Callable> wrappedSupplier = @@ -101,6 +103,56 @@ protected StoreMetaValue getStoreMetaValue(String storeName, StoreMetaKey key) { return value; } + @Override + protected SchemaData getSchemaData(String storeName) { + SchemaData schemaData = schemaMap.get(storeName); + SchemaEntry keySchema; + if (schemaData == null) { + // Retrieve the key schema and initialize SchemaData only if it's not cached yet. + StoreMetaKey keySchemaKey = MetaStoreDataType.STORE_KEY_SCHEMAS + .getStoreMetaKey(Collections.singletonMap(KEY_STRING_STORE_NAME, storeName)); + Map keySchemaMap = + getStoreMetaValue(storeName, keySchemaKey).storeKeySchemas.keySchemaMap; + if (keySchemaMap.isEmpty()) { + throw new VeniceException("No key schema found for store: " + storeName); + } + Map.Entry keySchemaEntry = keySchemaMap.entrySet().iterator().next(); + keySchema = + new SchemaEntry(Integer.parseInt(keySchemaEntry.getKey().toString()), keySchemaEntry.getValue().toString()); + schemaData = new SchemaData(storeName, keySchema); + } + StoreMetaKey valueSchemaKey = MetaStoreDataType.STORE_VALUE_SCHEMAS + .getStoreMetaKey(Collections.singletonMap(KEY_STRING_STORE_NAME, storeName)); + Map valueSchemaMap = + getStoreMetaValue(storeName, valueSchemaKey).storeValueSchemas.valueSchemaMap; + // Check the value schema string, if it's empty then try to query the other key space for individual value schema. + for (Map.Entry entry: valueSchemaMap.entrySet()) { + // Check if we already have the corresponding value schema + int valueSchemaId = Integer.parseInt(entry.getKey().toString()); + if (schemaData.getValueSchema(valueSchemaId) != null) { + continue; + } + if (entry.getValue().toString().isEmpty()) { + // The value schemas might be too large to be stored in a single K/V. + StoreMetaKey individualValueSchemaKey = + MetaStoreDataType.STORE_VALUE_SCHEMA.getStoreMetaKey(new HashMap() { + { + put(KEY_STRING_STORE_NAME, storeName); + put(KEY_STRING_SCHEMA_ID, entry.getKey().toString()); + } + }); + // Empty string is not a valid value schema therefore it's safe to throw exceptions if we also cannot find it in + // the individual value schema key space. + String valueSchema = + getStoreMetaValue(storeName, individualValueSchemaKey).storeValueSchema.valueSchema.toString(); + schemaData.addValueSchema(new SchemaEntry(valueSchemaId, valueSchema)); + } else { + schemaData.addValueSchema(new SchemaEntry(valueSchemaId, entry.getValue().toString())); + } + } + return schemaData; + } + private AvroSpecificStoreClient getAvroClientForMetaStore(String storeName) { return storeClientMap.computeIfAbsent(storeName, k -> { ClientConfig clonedClientConfig = ClientConfig.cloneConfig(clientConfig) @@ -112,4 +164,13 @@ private AvroSpecificStoreClient getAvroClientForMe return ClientFactory.getAndStartSpecificAvroClient(clonedClientConfig); }); } + + // Helper function with common code for retrieving StoreConfig from meta system store. + protected StoreConfig getStoreConfigFromMetaSystemStore(String storeName) { + StoreClusterConfig clusterConfig = getStoreMetaValue( + storeName, + MetaStoreDataType.STORE_CLUSTER_CONFIG + .getStoreMetaKey(Collections.singletonMap(KEY_STRING_STORE_NAME, storeName))).storeClusterConfig; + return new StoreConfig(clusterConfig); + } } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/ingestion/heartbeat/AggregatedHeartbeatLagEntry.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/ingestion/heartbeat/AggregatedHeartbeatLagEntry.java new file mode 100644 index 00000000000..01c8a34b64d --- /dev/null +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/ingestion/heartbeat/AggregatedHeartbeatLagEntry.java @@ -0,0 +1,19 @@ +package com.linkedin.davinci.stats.ingestion.heartbeat; + +public class AggregatedHeartbeatLagEntry { + private final long currentVersionHeartbeatLag; + private final long nonCurrentVersionHeartbeatLag; + + public AggregatedHeartbeatLagEntry(long currentVersionHeartbeatLag, long nonCurrentVersionHeartbeatLag) { + this.currentVersionHeartbeatLag = currentVersionHeartbeatLag; + this.nonCurrentVersionHeartbeatLag = nonCurrentVersionHeartbeatLag; + } + + public long getCurrentVersionHeartbeatLag() { + return currentVersionHeartbeatLag; + } + + public long getNonCurrentVersionHeartbeatLag() { + return nonCurrentVersionHeartbeatLag; + } +} diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/ingestion/heartbeat/HeartbeatMonitoringService.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/ingestion/heartbeat/HeartbeatMonitoringService.java index 5476dbbf8f5..fdd22b2ab04 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/ingestion/heartbeat/HeartbeatMonitoringService.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/stats/ingestion/heartbeat/HeartbeatMonitoringService.java @@ -3,6 +3,7 @@ import com.linkedin.davinci.kafka.consumer.LeaderFollowerStateType; import com.linkedin.davinci.kafka.consumer.ReplicaHeartbeatInfo; import com.linkedin.venice.meta.ReadOnlyStoreRepository; +import com.linkedin.venice.meta.Store; import com.linkedin.venice.meta.Version; import com.linkedin.venice.service.AbstractVeniceService; import com.linkedin.venice.utils.Utils; @@ -43,6 +44,7 @@ public class HeartbeatMonitoringService extends AbstractVeniceService { private static final Logger LOGGER = LogManager.getLogger(HeartbeatMonitoringService.class); + private final ReadOnlyStoreRepository metadataRepository; private final Thread reportingThread; private final Thread lagLoggingThread; @@ -65,6 +67,7 @@ public HeartbeatMonitoringService( this.lagLoggingThread = new HeartbeatLagLoggingThread(); this.followerHeartbeatTimeStamps = new VeniceConcurrentHashMap<>(); this.leaderHeartbeatTimeStamps = new VeniceConcurrentHashMap<>(); + this.metadataRepository = metadataRepository; this.versionStatsReporter = new HeartbeatVersionedStats( metricsRepository, metadataRepository, @@ -377,6 +380,48 @@ protected void checkAndMaybeLogHeartbeatDelay() { checkAndMaybeLogHeartbeatDelayMap(followerHeartbeatTimeStamps); } + AggregatedHeartbeatLagEntry getMaxHeartbeatLag( + Map>>> heartbeatTimestamps, + boolean isLeaderLag) { + long currentTimestamp = System.currentTimeMillis(); + long minHeartbeatTimestampForCurrentVersion = Long.MAX_VALUE; + long minHeartbeatTimestampForNonCurrentVersion = Long.MAX_VALUE; + for (Map.Entry>>> storeName: heartbeatTimestamps + .entrySet()) { + Store store = metadataRepository.getStore(storeName.getKey()); + if (store == null) { + LOGGER.warn("Store: {} not found in repository", storeName.getKey()); + continue; + } + int currentVersion = store.getCurrentVersion(); + for (Map.Entry>> version: storeName.getValue() + .entrySet()) { + for (Map.Entry> partition: version.getValue().entrySet()) { + for (Map.Entry region: partition.getValue().entrySet()) { + long heartbeatTs = region.getValue().timestamp; + if (currentVersion == version.getKey()) { + minHeartbeatTimestampForCurrentVersion = Math.min(minHeartbeatTimestampForCurrentVersion, heartbeatTs); + } else { + minHeartbeatTimestampForNonCurrentVersion = + Math.min(minHeartbeatTimestampForNonCurrentVersion, heartbeatTs); + } + } + } + } + } + return new AggregatedHeartbeatLagEntry( + currentTimestamp - minHeartbeatTimestampForCurrentVersion, + currentTimestamp - minHeartbeatTimestampForNonCurrentVersion); + } + + public AggregatedHeartbeatLagEntry getMaxLeaderHeartbeatLag() { + return getMaxHeartbeatLag(leaderHeartbeatTimeStamps, true); + } + + public AggregatedHeartbeatLagEntry getMaxFollowerHeartbeatLag() { + return getMaxHeartbeatLag(leaderHeartbeatTimeStamps, false); + } + @FunctionalInterface interface ReportLagFunction { void apply(String storeName, int version, String region, long lag, boolean isReadyToServe); diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBServerConfig.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBServerConfig.java index 820c22aac7c..bdcb1eb3153 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBServerConfig.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBServerConfig.java @@ -98,6 +98,9 @@ public class RocksDBServerConfig { * Max memtable count per database; */ public static final String ROCKSDB_MAX_MEMTABLE_COUNT = "rocksdb.max.memtable.count"; + + public static final String ROCKSDB_MIN_WRITE_BUFFER_NUMBER_TO_MERGE = "rocksdb.min.write.buffer.number.to.merge"; + /** * Max total WAL log size per database; */ @@ -253,6 +256,8 @@ public class RocksDBServerConfig { private final long rocksDBMemtableSizeInBytes; private final int rocksDBMaxMemtableCount; + + private final int rocksDBMinWriteBufferNumberToMerge; private final long rocksDBMaxTotalWalSizeInBytes; private final long rocksDBMaxBytesForLevelBase; @@ -351,6 +356,7 @@ public RocksDBServerConfig(VeniceProperties props) { this.rocksDBMemtableSizeInBytes = props.getSizeInBytes(ROCKSDB_MEMTABLE_SIZE_IN_BYTES, 32 * 1024 * 1024L); // 32MB this.rocksDBMaxMemtableCount = props.getInt(ROCKSDB_MAX_MEMTABLE_COUNT, 2); + this.rocksDBMinWriteBufferNumberToMerge = props.getInt(ROCKSDB_MIN_WRITE_BUFFER_NUMBER_TO_MERGE, 1); /** * Default: 0 means letting RocksDB to decide the proper WAL size. * Here is the related docs in RocksDB C++ lib: @@ -542,6 +548,10 @@ public int getRocksDBMaxMemtableCount() { return rocksDBMaxMemtableCount; } + public int getRocksDBMinWriteBufferNumberToMerge() { + return rocksDBMinWriteBufferNumberToMerge; + } + public long getRocksDBMaxTotalWalSizeInBytes() { return rocksDBMaxTotalWalSizeInBytes; } diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartition.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartition.java index f3e35005b76..3ec74037c24 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartition.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartition.java @@ -371,6 +371,7 @@ protected Options getStoreOptions(StoragePartitionConfig storagePartitionConfig, options.setMaxOpenFiles(rocksDBServerConfig.getMaxOpenFiles()); options.setTargetFileSizeBase(rocksDBServerConfig.getTargetFileSizeInBytes()); options.setMaxFileOpeningThreads(rocksDBServerConfig.getMaxFileOpeningThreads()); + options.setMinWriteBufferNumberToMerge(rocksDBServerConfig.getRocksDBMinWriteBufferNumberToMerge()); /** * Disable the stat dump threads, which will create excessive threads, which will eventually crash diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/BlobSnapshotManagerTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/BlobSnapshotManagerTest.java index 546365db145..8c707c55108 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/BlobSnapshotManagerTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/BlobSnapshotManagerTest.java @@ -1,5 +1,6 @@ package com.linkedin.davinci.blobtransfer; +import static com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferTableFormat; import static org.mockito.Mockito.any; import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doReturn; @@ -54,8 +55,12 @@ public class BlobSnapshotManagerTest { new BlobTransferPartitionMetadata(); private static final String DB_DIR = BASE_PATH + "/" + STORE_NAME + "_v" + VERSION_ID + "/" + RocksDBUtils.getPartitionDbName(STORE_NAME + "_v" + VERSION_ID, PARTITION_ID); - private static final BlobTransferPayload blobTransferPayload = - new BlobTransferPayload(BASE_PATH, STORE_NAME, VERSION_ID, PARTITION_ID); + private static final BlobTransferPayload blobTransferPayload = new BlobTransferPayload( + BASE_PATH, + STORE_NAME, + VERSION_ID, + PARTITION_ID, + BlobTransferTableFormat.BLOCK_BASED_TABLE); @Test(timeOut = TIMEOUT) public void testHybridSnapshot() { diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestNettyP2PBlobTransferManager.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestNettyP2PBlobTransferManager.java index 03288d90316..856cf708a07 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestNettyP2PBlobTransferManager.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestNettyP2PBlobTransferManager.java @@ -1,5 +1,6 @@ package com.linkedin.davinci.blobtransfer; +import static com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferTableFormat; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.doReturn; @@ -12,6 +13,7 @@ import com.linkedin.davinci.storage.StorageMetadataService; import com.linkedin.venice.blobtransfer.BlobFinder; import com.linkedin.venice.blobtransfer.BlobPeersDiscoveryResponse; +import com.linkedin.venice.exceptions.VeniceBlobTransferFileNotFoundException; import com.linkedin.venice.exceptions.VenicePeersConnectionException; import com.linkedin.venice.exceptions.VenicePeersNotFoundException; import com.linkedin.venice.kafka.protocol.state.PartitionState; @@ -117,7 +119,7 @@ public void teardown() throws Exception { public void testFailedConnectPeer() { CompletionStage future = null; try { - future = client.get("remotehost123", "test_store", 1, 1); + future = client.get("remotehost123", "test_store", 1, 1, BlobTransferTableFormat.BLOCK_BASED_TABLE); } catch (Exception e) { Assert.assertTrue(e instanceof VenicePeersConnectionException); Assert.assertEquals(e.getMessage(), "Failed to connect to the host: remotehost123"); @@ -132,7 +134,8 @@ public void testFailedConnectPeer() { @Test public void testFailedRequestFromFinder() { doReturn(null).when(finder).discoverBlobPeers(anyString(), anyInt(), anyInt()); - CompletionStage future = manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION); + CompletionStage future = + manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); Assert.assertTrue(future.toCompletableFuture().isCompletedExceptionally()); future.whenComplete((result, throwable) -> { Assert.assertNotNull(throwable); @@ -156,7 +159,8 @@ public void testNoResultFromFinder() { Mockito.doReturn(expectOffsetRecord).when(storageMetadataService).getLastOffset(Mockito.any(), Mockito.anyInt()); // Execution: - CompletionStage future = manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION); + CompletionStage future = + manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); Assert.assertTrue(future.toCompletableFuture().isCompletedExceptionally()); future.whenComplete((result, throwable) -> { Assert.assertNotNull(throwable); @@ -166,6 +170,42 @@ public void testNoResultFromFinder() { verifyFileTransferFailed(expectOffsetRecord); } + /** + * Testing the request blob transfer table format is not match to the existing snapshot format + * @throws IOException + */ + @Test + public void testSnapshotFormatNotMatch() throws IOException { + // Preparation: + Mockito.doReturn(false).when(blobSnapshotManager).isStoreHybrid(anyString(), anyInt()); + + BlobPeersDiscoveryResponse response = new BlobPeersDiscoveryResponse(); + response.setDiscoveryResult(Collections.singletonList("localhost")); + doReturn(response).when(finder).discoverBlobPeers(anyString(), anyInt(), anyInt()); + + StoreVersionState storeVersionState = new StoreVersionState(); + Mockito.doReturn(storeVersionState).when(storageMetadataService).getStoreVersionState(Mockito.any()); + + InternalAvroSpecificSerializer partitionStateSerializer = + AvroProtocolDefinition.PARTITION_STATE.getSerializer(); + OffsetRecord expectOffsetRecord = new OffsetRecord(partitionStateSerializer); + expectOffsetRecord.setOffsetLag(1000L); + Mockito.doReturn(expectOffsetRecord).when(storageMetadataService).getLastOffset(Mockito.any(), Mockito.anyInt()); + + snapshotPreparation(); + + // Execution: + // Bootstrap try to get plain table snapshot, + // but the existing snapshot is block based table, as the snapshot manager is initialized with block based table + // config + CompletionStage future = + manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.PLAIN_TABLE); + future.whenComplete((result, throwable) -> { + Assert.assertNotNull(throwable); + Assert.assertTrue(throwable instanceof VeniceBlobTransferFileNotFoundException); + }); + } + @Test public void testLocalFileTransferInBatchStore() throws IOException, ExecutionException, InterruptedException, TimeoutException { @@ -189,7 +229,8 @@ public void testLocalFileTransferInBatchStore() // Execution: // Manager should be able to fetch the file and download it to another directory - CompletionStage future = manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION); + CompletionStage future = + manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); future.toCompletableFuture().get(1, TimeUnit.MINUTES); // Verification: @@ -224,7 +265,8 @@ public void testSkipBadHostAndUseCorrectHost() // Execution: // Manager should be able to fetch the file and download it to another directory, and future is done normally - CompletionStage future = manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION); + CompletionStage future = + manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); future.toCompletableFuture().get(1, TimeUnit.MINUTES); future.whenComplete((result, throwable) -> { Assert.assertNull(throwable); @@ -232,9 +274,12 @@ public void testSkipBadHostAndUseCorrectHost() // Verification: // Verify that even has bad hosts in the list, it still finally uses good host to transfer the file - Mockito.verify(client, Mockito.times(1)).get("localhost", TEST_STORE, TEST_VERSION, TEST_PARTITION); - Mockito.verify(client, Mockito.times(1)).get("badhost1", TEST_STORE, TEST_VERSION, TEST_PARTITION); - Mockito.verify(client, Mockito.times(1)).get("badhost2", TEST_STORE, TEST_VERSION, TEST_PARTITION); + Mockito.verify(client, Mockito.times(1)) + .get("localhost", TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); + Mockito.verify(client, Mockito.times(1)) + .get("badhost1", TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); + Mockito.verify(client, Mockito.times(1)) + .get("badhost2", TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); verifyFileTransferSuccess(expectOffsetRecord); } @@ -266,7 +311,8 @@ public void testUseCorrectHostAndSkipRemainingHosts() // Execution: // Manager should be able to fetch the file and download it to another directory, and future is done normally - CompletionStage future = manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION); + CompletionStage future = + manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); future.toCompletableFuture().get(1, TimeUnit.MINUTES); future.whenComplete((result, throwable) -> { Assert.assertNull(throwable); @@ -275,9 +321,12 @@ public void testUseCorrectHostAndSkipRemainingHosts() // Verification: // Verify that it has the localhost in the first place, it should use localhost to transfer the file // All the remaining bad hosts should not be called to fetch the file. - Mockito.verify(client, Mockito.times(1)).get("localhost", TEST_STORE, TEST_VERSION, TEST_PARTITION); - Mockito.verify(client, Mockito.never()).get("badhost1", TEST_STORE, TEST_VERSION, TEST_PARTITION); - Mockito.verify(client, Mockito.never()).get("badhost2", TEST_STORE, TEST_VERSION, TEST_PARTITION); + Mockito.verify(client, Mockito.times(1)) + .get("localhost", TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); + Mockito.verify(client, Mockito.never()) + .get("badhost1", TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); + Mockito.verify(client, Mockito.never()) + .get("badhost2", TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); verifyFileTransferSuccess(expectOffsetRecord); } @@ -306,7 +355,8 @@ public void testLocalFileTransferInHybridStore() // Execution: // Manager should be able to fetch the file and download it to another directory - CompletionStage future = manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION); + CompletionStage future = + manager.get(TEST_STORE, TEST_VERSION, TEST_PARTITION, BlobTransferTableFormat.BLOCK_BASED_TABLE); future.toCompletableFuture().get(1, TimeUnit.MINUTES); // Verification: @@ -394,7 +444,7 @@ private void verifyFileTransferSuccess(OffsetRecord expectOffsetRecord) throws I */ private void verifyFileTransferFailed(OffsetRecord expectOffsetRecord) { // Verify client never get called - Mockito.verify(client, Mockito.never()).get(anyString(), anyString(), anyInt(), anyInt()); + Mockito.verify(client, Mockito.never()).get(anyString(), anyString(), anyInt(), anyInt(), Mockito.any()); // Verify files are not written to the partition directory Assert.assertFalse(Files.exists(destFile1)); diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestP2PFileTransferClientHandler.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestP2PFileTransferClientHandler.java index 803a18881cf..f5f4eee49b4 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestP2PFileTransferClientHandler.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestP2PFileTransferClientHandler.java @@ -72,7 +72,8 @@ public void setUp() throws IOException { inputStreamFuture, TEST_STORE, TEST_VERSION, - TEST_PARTITION)); + TEST_PARTITION, + BlobTransferUtils.BlobTransferTableFormat.BLOCK_BASED_TABLE)); clientMetadataHandler = Mockito.spy( new P2PMetadataTransferHandler( @@ -80,7 +81,8 @@ public void setUp() throws IOException { baseDir.toString(), TEST_STORE, TEST_VERSION, - TEST_PARTITION)); + TEST_PARTITION, + BlobTransferUtils.BlobTransferTableFormat.BLOCK_BASED_TABLE)); Mockito.doNothing().when(clientMetadataHandler).updateStorePartitionMetadata(Mockito.any(), Mockito.any()); @@ -216,7 +218,12 @@ public void testSingleFileTransfer() throws ExecutionException, InterruptedExcep inputStreamFuture.toCompletableFuture().get(1, TimeUnit.MINUTES); // verify the content is written to the disk - BlobTransferPayload payload = new BlobTransferPayload(baseDir.toString(), TEST_STORE, TEST_VERSION, TEST_PARTITION); + BlobTransferPayload payload = new BlobTransferPayload( + baseDir.toString(), + TEST_STORE, + TEST_VERSION, + TEST_PARTITION, + BlobTransferUtils.BlobTransferTableFormat.BLOCK_BASED_TABLE); Path dest = Paths.get(payload.getPartitionDir()); Assert.assertTrue(Files.exists(dest)); Assert.assertTrue(Files.isDirectory(dest)); @@ -259,7 +266,12 @@ public void testMultipleFilesTransfer() inputStreamFuture.toCompletableFuture().get(1, TimeUnit.MINUTES); // verify the content is written to the disk - BlobTransferPayload payload = new BlobTransferPayload(baseDir.toString(), TEST_STORE, TEST_VERSION, TEST_PARTITION); + BlobTransferPayload payload = new BlobTransferPayload( + baseDir.toString(), + TEST_STORE, + TEST_VERSION, + TEST_PARTITION, + BlobTransferUtils.BlobTransferTableFormat.BLOCK_BASED_TABLE); Path dest = Paths.get(payload.getPartitionDir()); Assert.assertTrue(Files.exists(dest)); Assert.assertTrue(Files.isDirectory(dest)); @@ -381,7 +393,12 @@ public void testMultipleFilesAndOneMetadataTransfer() inputStreamFuture.toCompletableFuture().get(1, TimeUnit.MINUTES); // Verify the files are written to disk - BlobTransferPayload payload = new BlobTransferPayload(baseDir.toString(), TEST_STORE, TEST_VERSION, TEST_PARTITION); + BlobTransferPayload payload = new BlobTransferPayload( + baseDir.toString(), + TEST_STORE, + TEST_VERSION, + TEST_PARTITION, + BlobTransferUtils.BlobTransferTableFormat.BLOCK_BASED_TABLE); Path dest = Paths.get(payload.getPartitionDir()); Assert.assertTrue(Files.exists(dest)); Assert.assertTrue(Files.isDirectory(dest)); diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestP2PFileTransferServerHandler.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestP2PFileTransferServerHandler.java index 6197269ba8b..2fa4343f2bd 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestP2PFileTransferServerHandler.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/blobtransfer/TestP2PFileTransferServerHandler.java @@ -90,12 +90,15 @@ public void testRejectNonGETMethod() { Assert.assertEquals(response.status().code(), 405); } + /** + * Testing the method is GET, but uri format is invalid + */ @Test public void testRejectInvalidPath() { FullHttpRequest request = new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/test"); ch.writeInbound(request); FullHttpResponse response = ch.readOutbound(); - Assert.assertEquals(response.status().code(), 400); + Assert.assertEquals(response.status().code(), 500); } @Test @@ -109,7 +112,36 @@ public void testRejectNonExistPath() { offsetRecord.setOffsetLag(1000L); Mockito.doReturn(offsetRecord).when(storageMetadataService).getLastOffset(Mockito.any(), Mockito.anyInt()); - FullHttpRequest request = new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/myStore/1/10"); + FullHttpRequest request = + new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/myStore/1/10/BLOCK_BASED_TABLE"); + ch.writeInbound(request); + FullHttpResponse response = ch.readOutbound(); + Assert.assertEquals(response.status().code(), 404); + Assert.assertEquals(blobSnapshotManager.getConcurrentSnapshotUsers("myStore_v1", 10), 0); + } + + /** + * Sending request for plain table format, but the snapshot manager is use the block based format. + */ + @Test + public void testRejectNotMatchFormat() throws IOException { + // prepare the file request + Path snapshotDir = Paths.get(RocksDBUtils.composeSnapshotDir(baseDir.toString(), "myStore_v1", 10)); + Files.createDirectories(snapshotDir); + Path file1 = snapshotDir.resolve("file1"); + Files.write(file1.toAbsolutePath(), "hello".getBytes()); + + // prepare response from metadata service for the metadata preparation + StoreVersionState storeVersionState = new StoreVersionState(); + Mockito.doReturn(storeVersionState).when(storageMetadataService).getStoreVersionState(Mockito.any()); + InternalAvroSpecificSerializer partitionStateSerializer = + AvroProtocolDefinition.PARTITION_STATE.getSerializer(); + OffsetRecord offsetRecord = new OffsetRecord(partitionStateSerializer); + offsetRecord.setOffsetLag(1000L); + Mockito.doReturn(offsetRecord).when(storageMetadataService).getLastOffset(Mockito.any(), Mockito.anyInt()); + + FullHttpRequest request = + new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/myStore/1/10/PLAIN_TABLE"); ch.writeInbound(request); FullHttpResponse response = ch.readOutbound(); Assert.assertEquals(response.status().code(), 404); @@ -130,7 +162,8 @@ public void testFailOnAccessPath() throws IOException { // create an empty snapshot dir Path snapshotDir = Paths.get(RocksDBUtils.composeSnapshotDir(baseDir.toString(), "myStore_v1", 10)); Files.createDirectories(snapshotDir); - FullHttpRequest request = new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/myStore/1/10"); + FullHttpRequest request = + new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/myStore/1/10/BLOCK_BASED_TABLE"); ch.writeInbound(request); FullHttpResponse response = ch.readOutbound(); @@ -164,7 +197,8 @@ public void testTransferSingleFileAndSingleMetadataForBatchStore() throws IOExce Path file1 = snapshotDir.resolve("file1"); Files.write(file1.toAbsolutePath(), "hello".getBytes()); String file1ChecksumHeader = BlobTransferUtils.generateFileChecksum(file1); - FullHttpRequest request = new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/myStore/1/10"); + FullHttpRequest request = + new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/myStore/1/10/BLOCK_BASED_TABLE"); ch.writeInbound(request); @@ -225,7 +259,8 @@ public void testTransferMultipleFiles() throws IOException { Files.write(file2.toAbsolutePath(), "world".getBytes()); String file2ChecksumHeader = BlobTransferUtils.generateFileChecksum(file2); - FullHttpRequest request = new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/myStore/1/10"); + FullHttpRequest request = + new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/myStore/1/10/BLOCK_BASED_TABLE"); Set fileNames = new HashSet<>(); Set fileChecksums = new HashSet<>(); // the order of file transfer is not guaranteed so put them into a set and remove them one by one @@ -303,7 +338,8 @@ public void testWhenMetadataCreateError() throws IOException { Files.createDirectories(snapshotDir); Path file1 = snapshotDir.resolve("file1"); Files.write(file1.toAbsolutePath(), "hello".getBytes()); - FullHttpRequest request = new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/myStore/1/10"); + FullHttpRequest request = + new DefaultFullHttpRequest(HttpVersion.HTTP_1_1, HttpMethod.GET, "/myStore/1/10/BLOCK_BASED_TABLE"); ch.writeInbound(request); diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/client/AvroGenericDaVinciClientTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/client/AvroGenericDaVinciClientTest.java index cf84fd4a7a1..f38915ee329 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/client/AvroGenericDaVinciClientTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/client/AvroGenericDaVinciClientTest.java @@ -61,7 +61,12 @@ public AvroGenericDaVinciClient setUpClientWithRecordTransformer( } DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( - (storeVersion) -> new TestStringRecordTransformer(storeVersion, true), + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> new TestStringRecordTransformer( + storeVersion, + keySchema, + inputValueSchema, + outputValueSchema, + true), String.class, Schema.create(Schema.Type.STRING)); daVinciConfig.setRecordTransformerConfig(recordTransformerConfig); diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/config/DaVinciConfigTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/config/DaVinciConfigTest.java index 70e7a9c93be..fab664eed59 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/config/DaVinciConfigTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/config/DaVinciConfigTest.java @@ -16,18 +16,13 @@ public class DaVinciConfigTest { public class TestRecordTransformer extends DaVinciRecordTransformer { - public TestRecordTransformer(int storeVersion, boolean storeRecordsInDaVinci) { - super(storeVersion, storeRecordsInDaVinci); - } - - @Override - public Schema getKeySchema() { - return Schema.create(Schema.Type.INT); - } - - @Override - public Schema getOutputValueSchema() { - return Schema.create(Schema.Type.INT); + public TestRecordTransformer( + int storeVersion, + Schema keySchema, + Schema inputValueSchema, + Schema outputValueSchema, + boolean storeRecordsInDaVinci) { + super(storeVersion, keySchema, inputValueSchema, outputValueSchema, storeRecordsInDaVinci); } @Override @@ -45,9 +40,15 @@ public void processPut(Lazy key, Lazy value) { public void testRecordTransformerEnabled() { DaVinciConfig config = new DaVinciConfig(); assertFalse(config.isRecordTransformerEnabled()); + DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( - (storeVersion) -> new TestRecordTransformer(storeVersion, true), - Integer.class, + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> new TestRecordTransformer( + storeVersion, + keySchema, + inputValueSchema, + outputValueSchema, + true), + String.class, Schema.create(Schema.Type.INT)); config.setRecordTransformerConfig(recordTransformerConfig); assertTrue(config.isRecordTransformerEnabled()); @@ -55,15 +56,20 @@ public void testRecordTransformerEnabled() { @Test public void testGetAndSetRecordTransformer() { - Integer testStoreVersion = 1; DaVinciConfig config = new DaVinciConfig(); - assertNull(config.getRecordTransformer(testStoreVersion)); + assertNull(config.getRecordTransformerConfig()); + DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( - (storeVersion) -> new TestRecordTransformer(storeVersion, true), - Integer.class, + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> new TestRecordTransformer( + storeVersion, + keySchema, + inputValueSchema, + outputValueSchema, + true), + String.class, Schema.create(Schema.Type.INT)); config.setRecordTransformerConfig(recordTransformerConfig); - assertNotNull(config.getRecordTransformer(testStoreVersion)); + assertNotNull(config.getRecordTransformerConfig()); } } diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/consumer/VeniceChangelogConsumerClientFactoryTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/consumer/VeniceChangelogConsumerClientFactoryTest.java index 6b0136f4764..ae02c6a5113 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/consumer/VeniceChangelogConsumerClientFactoryTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/consumer/VeniceChangelogConsumerClientFactoryTest.java @@ -3,7 +3,8 @@ import static com.linkedin.venice.ConfigKeys.CLUSTER_NAME; import static com.linkedin.venice.ConfigKeys.KAFKA_BOOTSTRAP_SERVERS; import static com.linkedin.venice.ConfigKeys.ZOOKEEPER_ADDRESS; -import static org.mockito.Mockito.*; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; import com.fasterxml.jackson.core.JsonProcessingException; import com.linkedin.d2.balancer.D2Client; diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/ingestion/DefaultIngestionBackendTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/ingestion/DefaultIngestionBackendTest.java index fe6910cac66..82dc5468235 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/ingestion/DefaultIngestionBackendTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/ingestion/DefaultIngestionBackendTest.java @@ -1,5 +1,6 @@ package com.linkedin.davinci.ingestion; +import static com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferTableFormat; import static org.mockito.ArgumentMatchers.anyInt; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.any; @@ -21,6 +22,7 @@ import com.linkedin.davinci.storage.StorageMetadataService; import com.linkedin.davinci.storage.StorageService; import com.linkedin.davinci.store.AbstractStorageEngine; +import com.linkedin.davinci.store.rocksdb.RocksDBServerConfig; import com.linkedin.venice.exceptions.VenicePeersNotFoundException; import com.linkedin.venice.kafka.protocol.state.StoreVersionState; import com.linkedin.venice.meta.ReadOnlyStoreRepository; @@ -32,6 +34,7 @@ import java.time.Duration; import java.util.concurrent.CompletableFuture; import org.mockito.Mock; +import org.mockito.Mockito; import org.mockito.MockitoAnnotations; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -72,6 +75,7 @@ public class DefaultIngestionBackendTest { private static final String STORE_NAME = "testStore"; private static final String STORE_VERSION = "store_v1"; private static final String BASE_DIR = "mockBaseDir"; + private static final BlobTransferTableFormat BLOB_TRANSFER_FORMAT = BlobTransferTableFormat.BLOCK_BASED_TABLE; @BeforeMethod public void setUp() { @@ -108,12 +112,15 @@ public void setUp() { public void testStartConsumptionWithBlobTransfer() { when(store.isBlobTransferEnabled()).thenReturn(true); when(store.isHybrid()).thenReturn(true); - when(blobTransferManager.get(eq(STORE_NAME), eq(VERSION_NUMBER), eq(PARTITION))) + when(blobTransferManager.get(eq(STORE_NAME), eq(VERSION_NUMBER), eq(PARTITION), eq(BLOB_TRANSFER_FORMAT))) .thenReturn(CompletableFuture.completedFuture(null)); when(veniceServerConfig.getRocksDBPath()).thenReturn(BASE_DIR); + RocksDBServerConfig rocksDBServerConfig = Mockito.mock(RocksDBServerConfig.class); + when(rocksDBServerConfig.isRocksDBPlainTableFormatEnabled()).thenReturn(false); + when(veniceServerConfig.getRocksDBServerConfig()).thenReturn(rocksDBServerConfig); ingestionBackend.startConsumption(storeConfig, PARTITION); - verify(blobTransferManager).get(eq(STORE_NAME), eq(VERSION_NUMBER), eq(PARTITION)); + verify(blobTransferManager).get(eq(STORE_NAME), eq(VERSION_NUMBER), eq(PARTITION), eq(BLOB_TRANSFER_FORMAT)); verify(aggVersionedBlobTransferStats).recordBlobTransferResponsesCount(eq(STORE_NAME), eq(VERSION_NUMBER)); verify(aggVersionedBlobTransferStats) .recordBlobTransferResponsesBasedOnBoostrapStatus(eq(STORE_NAME), eq(VERSION_NUMBER), eq(true)); @@ -125,10 +132,12 @@ public void testStartConsumptionWithBlobTransferWhenNoPeerFound() { when(store.isHybrid()).thenReturn(false); CompletableFuture errorFuture = new CompletableFuture<>(); errorFuture.completeExceptionally(new VenicePeersNotFoundException("No peers found")); - when(blobTransferManager.get(eq(STORE_NAME), eq(VERSION_NUMBER), eq(PARTITION))).thenReturn(errorFuture); + when(blobTransferManager.get(eq(STORE_NAME), eq(VERSION_NUMBER), eq(PARTITION), eq(BLOB_TRANSFER_FORMAT))) + .thenReturn(errorFuture); CompletableFuture future = - ingestionBackend.bootstrapFromBlobs(store, VERSION_NUMBER, PARTITION, 100L).toCompletableFuture(); + ingestionBackend.bootstrapFromBlobs(store, VERSION_NUMBER, PARTITION, BLOB_TRANSFER_FORMAT, 100L) + .toCompletableFuture(); assertTrue(future.isDone()); verify(aggVersionedBlobTransferStats).recordBlobTransferResponsesCount(eq(STORE_NAME), eq(VERSION_NUMBER)); verify(aggVersionedBlobTransferStats) @@ -146,12 +155,15 @@ public void testNotStartBootstrapFromBlobTransferWhenNotLagging() { when(store.isBlobTransferEnabled()).thenReturn(true); when(store.isHybrid()).thenReturn(false); CompletableFuture future = new CompletableFuture<>(); - when(blobTransferManager.get(eq(STORE_NAME), eq(VERSION_NUMBER), eq(PARTITION))).thenReturn(future); + when(blobTransferManager.get(eq(STORE_NAME), eq(VERSION_NUMBER), eq(PARTITION), eq(BLOB_TRANSFER_FORMAT))) + .thenReturn(future); CompletableFuture result = - ingestionBackend.bootstrapFromBlobs(store, VERSION_NUMBER, PARTITION, laggingThreshold).toCompletableFuture(); + ingestionBackend.bootstrapFromBlobs(store, VERSION_NUMBER, PARTITION, BLOB_TRANSFER_FORMAT, laggingThreshold) + .toCompletableFuture(); assertTrue(result.isDone()); - verify(blobTransferManager, never()).get(eq(STORE_NAME), eq(VERSION_NUMBER), eq(PARTITION)); + verify(blobTransferManager, never()) + .get(eq(STORE_NAME), eq(VERSION_NUMBER), eq(PARTITION), eq(BLOB_TRANSFER_FORMAT)); verify(aggVersionedBlobTransferStats, never()).recordBlobTransferResponsesCount(eq(STORE_NAME), eq(VERSION_NUMBER)); verify(aggVersionedBlobTransferStats, never()) .recordBlobTransferResponsesBasedOnBoostrapStatus(eq(STORE_NAME), eq(VERSION_NUMBER), eq(false)); diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/AdaptiveThrottlerSingalServiceTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/AdaptiveThrottlerSingalServiceTest.java new file mode 100644 index 00000000000..b42c2b723eb --- /dev/null +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/AdaptiveThrottlerSingalServiceTest.java @@ -0,0 +1,73 @@ +package com.linkedin.davinci.kafka.consumer; + +import static com.linkedin.davinci.kafka.consumer.AdaptiveThrottlerSignalService.SINGLE_GET_LATENCY_P99_METRIC_NAME; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.when; + +import com.linkedin.davinci.config.VeniceServerConfig; +import com.linkedin.davinci.stats.ingestion.heartbeat.AggregatedHeartbeatLagEntry; +import com.linkedin.davinci.stats.ingestion.heartbeat.HeartbeatMonitoringService; +import io.tehuti.Metric; +import io.tehuti.metrics.MetricsRepository; +import java.util.concurrent.TimeUnit; +import org.mockito.Mockito; +import org.testng.Assert; +import org.testng.annotations.Test; + + +public class AdaptiveThrottlerSingalServiceTest { + @Test + public void testUpdateSignal() { + MetricsRepository metricsRepository = mock(MetricsRepository.class); + HeartbeatMonitoringService heartbeatMonitoringService = mock(HeartbeatMonitoringService.class); + VeniceServerConfig veniceServerConfig = mock(VeniceServerConfig.class); + when(veniceServerConfig.getAdaptiveThrottlerSingleGetLatencyThreshold()).thenReturn(10d); + AdaptiveThrottlerSignalService adaptiveThrottlerSignalService = + new AdaptiveThrottlerSignalService(veniceServerConfig, metricsRepository, heartbeatMonitoringService); + + // Single Get Signal + Assert.assertFalse(adaptiveThrottlerSignalService.isSingleGetLatencySignalActive()); + Metric singleGetMetric = mock(Metric.class); + when(singleGetMetric.value()).thenReturn(20.0d); + when(metricsRepository.getMetric(SINGLE_GET_LATENCY_P99_METRIC_NAME)).thenReturn(singleGetMetric); + adaptiveThrottlerSignalService.refreshSignalAndThrottler(); + Assert.assertTrue(adaptiveThrottlerSignalService.isSingleGetLatencySignalActive()); + when(singleGetMetric.value()).thenReturn(1.0d); + Assert.assertTrue(adaptiveThrottlerSignalService.isSingleGetLatencySignalActive()); + adaptiveThrottlerSignalService.refreshSignalAndThrottler(); + Assert.assertFalse(adaptiveThrottlerSignalService.isSingleGetLatencySignalActive()); + + // Heartbeat signal + Assert.assertFalse(adaptiveThrottlerSignalService.isCurrentFollowerMaxHeartbeatLagSignalActive()); + Assert.assertFalse(adaptiveThrottlerSignalService.isCurrentLeaderMaxHeartbeatLagSignalActive()); + Assert.assertFalse(adaptiveThrottlerSignalService.isNonCurrentFollowerMaxHeartbeatLagSignalActive()); + Assert.assertFalse(adaptiveThrottlerSignalService.isNonCurrentLeaderMaxHeartbeatLagSignalActive()); + + when(heartbeatMonitoringService.getMaxLeaderHeartbeatLag()) + .thenReturn(new AggregatedHeartbeatLagEntry(TimeUnit.MINUTES.toMillis(100), TimeUnit.MINUTES.toMillis(1))); + when(heartbeatMonitoringService.getMaxFollowerHeartbeatLag()) + .thenReturn(new AggregatedHeartbeatLagEntry(TimeUnit.MINUTES.toMillis(1), TimeUnit.MINUTES.toMillis(100))); + adaptiveThrottlerSignalService.refreshSignalAndThrottler(); + Assert.assertFalse(adaptiveThrottlerSignalService.isCurrentFollowerMaxHeartbeatLagSignalActive()); + Assert.assertTrue(adaptiveThrottlerSignalService.isCurrentLeaderMaxHeartbeatLagSignalActive()); + Assert.assertTrue(adaptiveThrottlerSignalService.isNonCurrentFollowerMaxHeartbeatLagSignalActive()); + Assert.assertFalse(adaptiveThrottlerSignalService.isNonCurrentLeaderMaxHeartbeatLagSignalActive()); + } + + @Test + public void testRegisterThrottler() { + MetricsRepository metricsRepository = mock(MetricsRepository.class); + HeartbeatMonitoringService heartbeatMonitoringService = mock(HeartbeatMonitoringService.class); + VeniceServerConfig veniceServerConfig = mock(VeniceServerConfig.class); + when(veniceServerConfig.getAdaptiveThrottlerSingleGetLatencyThreshold()).thenReturn(10d); + AdaptiveThrottlerSignalService adaptiveThrottlerSignalService = + new AdaptiveThrottlerSignalService(veniceServerConfig, metricsRepository, heartbeatMonitoringService); + VeniceAdaptiveIngestionThrottler adaptiveIngestionThrottler = mock(VeniceAdaptiveIngestionThrottler.class); + adaptiveThrottlerSignalService.registerThrottler(adaptiveIngestionThrottler); + Assert.assertEquals(adaptiveThrottlerSignalService.getThrottlerList().size(), 1); + Assert.assertEquals(adaptiveThrottlerSignalService.getThrottlerList().get(0), adaptiveIngestionThrottler); + adaptiveThrottlerSignalService.refreshSignalAndThrottler(); + Mockito.verify(adaptiveIngestionThrottler, times(1)).checkSignalAndAdjustThrottler(); + } +} diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/IngestionThrottlerTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/IngestionThrottlerTest.java index fbd9a7c773a..1af79827d3f 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/IngestionThrottlerTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/IngestionThrottlerTest.java @@ -42,7 +42,8 @@ public void throttlerSwitchTest() throws IOException { ConcurrentHashMap tasks = new ConcurrentHashMap<>(); tasks.put("non_current_version_task", nonCurrentVersionTask); - IngestionThrottler throttler = new IngestionThrottler(true, serverConfig, () -> tasks, 10, TimeUnit.MILLISECONDS); + IngestionThrottler throttler = + new IngestionThrottler(true, serverConfig, () -> tasks, 10, TimeUnit.MILLISECONDS, null); TestUtils.waitForNonDeterministicAssertion(3, TimeUnit.SECONDS, () -> { assertFalse( throttler.isUsingSpeedupThrottler(), @@ -70,7 +71,7 @@ public void throttlerSwitchTest() throws IOException { tasks.clear(); IngestionThrottler throttlerForNonDaVinciClient = - new IngestionThrottler(false, serverConfig, () -> tasks, 10, TimeUnit.MILLISECONDS); + new IngestionThrottler(false, serverConfig, () -> tasks, 10, TimeUnit.MILLISECONDS, null); tasks.put("current_version_bootstrapping_task", currentVersionBootstrappingTask); tasks.put("current_version_completed_task", currentVersionCompletedTask); TestUtils.waitForNonDeterministicAssertion(3, TimeUnit.SECONDS, () -> { @@ -85,11 +86,11 @@ public void throttlerSwitchTest() throws IOException { @Test public void testDifferentThrottler() { VeniceServerConfig serverConfig = mock(VeniceServerConfig.class); - doReturn(100l).when(serverConfig).getKafkaFetchQuotaRecordPerSecond(); - doReturn(60l).when(serverConfig).getKafkaFetchQuotaTimeWindow(); - doReturn(1024l).when(serverConfig).getKafkaFetchQuotaBytesPerSecond(); + doReturn(100L).when(serverConfig).getKafkaFetchQuotaRecordPerSecond(); + doReturn(60L).when(serverConfig).getKafkaFetchQuotaTimeWindow(); + doReturn(1024L).when(serverConfig).getKafkaFetchQuotaBytesPerSecond(); IngestionThrottler ingestionThrottler = - new IngestionThrottler(true, serverConfig, () -> Collections.emptyMap(), 10, TimeUnit.MILLISECONDS); + new IngestionThrottler(true, serverConfig, () -> Collections.emptyMap(), 10, TimeUnit.MILLISECONDS, null); EventThrottler throttlerForAAWCLeader = mock(EventThrottler.class); EventThrottler throttlerForCurrentVersionAAWCLeader = mock(EventThrottler.class); EventThrottler throttlerForCurrentVersionNonAAWCLeader = mock(EventThrottler.class); diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionServiceTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionServiceTest.java index 21a79f82918..8555aa80a9b 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionServiceTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/KafkaStoreIngestionServiceTest.java @@ -181,6 +181,7 @@ public void testDisableMetricsEmission() { mockPubSubClientsFactory, Optional.empty(), null, + null, null); String mockStoreName = "test"; @@ -266,6 +267,7 @@ public void testGetIngestingTopicsNotWithOnlineVersion() { mockPubSubClientsFactory, Optional.empty(), null, + null, null); String topic1 = "test-store_v1"; String topic2 = "test-store_v2"; @@ -355,6 +357,7 @@ public void testCloseStoreIngestionTask() { mockPubSubClientsFactory, Optional.empty(), null, + null, null); String topicName = "test-store_v1"; String storeName = Version.parseStoreFromKafkaTopicName(topicName); @@ -421,6 +424,7 @@ public void testStoreIngestionTaskShutdownLastPartition(boolean isIsolatedIngest mockPubSubClientsFactory, Optional.empty(), null, + null, null); String topicName = "test-store_v1"; String storeName = Version.parseStoreFromKafkaTopicName(topicName); @@ -546,7 +550,7 @@ public void testDropStoragePartitionGracefully() throws NoSuchFieldException, Il storageServiceField.setAccessible(true); storageServiceField.set(storeIngestionTask, storageService); - Field storeConfigField = StoreIngestionTask.class.getDeclaredField("storeConfig"); + Field storeConfigField = StoreIngestionTask.class.getDeclaredField("storeVersionConfig"); storeConfigField.setAccessible(true); storeConfigField.set(storeIngestionTask, config); diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTaskTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTaskTest.java index f8285f64eab..9a4093012e8 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTaskTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/StoreIngestionTaskTest.java @@ -11,6 +11,7 @@ import static com.linkedin.davinci.kafka.consumer.StoreIngestionTaskTest.NodeType.LEADER; import static com.linkedin.davinci.kafka.consumer.StoreIngestionTaskTest.SortedInput.SORTED; import static com.linkedin.davinci.store.AbstractStorageEngine.StoragePartitionAdjustmentTrigger.PREPARE_FOR_READ; +import static com.linkedin.venice.ConfigKeys.BLOB_TRANSFER_MANAGER_ENABLED; import static com.linkedin.venice.ConfigKeys.CLUSTER_NAME; import static com.linkedin.venice.ConfigKeys.FREEZE_INGESTION_IF_READY_TO_SERVE_OR_LOCAL_DATA_EXISTS; import static com.linkedin.venice.ConfigKeys.HYBRID_QUOTA_ENFORCEMENT_ENABLED; @@ -79,7 +80,7 @@ import static org.testng.Assert.assertTrue; import static org.testng.Assert.fail; -import com.linkedin.davinci.client.DaVinciRecordTransformerFunctionalInterface; +import com.linkedin.davinci.client.DaVinciRecordTransformerConfig; import com.linkedin.davinci.compression.StorageEngineBackedCompressorFactory; import com.linkedin.davinci.config.VeniceServerConfig; import com.linkedin.davinci.config.VeniceStoreVersionConfig; @@ -348,6 +349,7 @@ public static Object[][] sortedInputAndAAConfigProvider() { private List mockNotifierError; private StorageMetadataService mockStorageMetadataService; private AbstractStorageEngine mockAbstractStorageEngine; + private DeepCopyStorageEngine mockDeepCopyStorageEngine; private IngestionThrottler mockIngestionThrottler; private Map kafkaUrlToRecordsThrottler; private KafkaClusterBasedRecordThrottler kafkaClusterBasedRecordThrottler; @@ -655,7 +657,7 @@ public static class StoreIngestionTaskTestConfig { private Optional diskUsageForTest = Optional.empty(); private Map extraServerProperties = new HashMap<>(); private Consumer storeVersionConfigOverride = storeVersionConfigOverride -> {}; - private DaVinciRecordTransformerFunctionalInterface recordTransformerFunction = null; + private DaVinciRecordTransformerConfig recordTransformerConfig = null; private OffsetRecord offsetRecord = null; public StoreIngestionTaskTestConfig(Set partitions, Runnable assertions, AAConfig aaConfig) { @@ -758,13 +760,13 @@ public StoreIngestionTaskTestConfig setStoreVersionConfigOverride( return this; } - public DaVinciRecordTransformerFunctionalInterface getRecordTransformerFunction() { - return recordTransformerFunction; + public DaVinciRecordTransformerConfig getRecordTransformerConfig() { + return recordTransformerConfig; } - public StoreIngestionTaskTestConfig setRecordTransformerFunction( - DaVinciRecordTransformerFunctionalInterface recordTransformerFunction) { - this.recordTransformerFunction = recordTransformerFunction; + public StoreIngestionTaskTestConfig setRecordTransformerConfig( + DaVinciRecordTransformerConfig recordTransformerConfig) { + this.recordTransformerConfig = recordTransformerConfig; return this; } @@ -792,7 +794,7 @@ private void runTest(StoreIngestionTaskTestConfig config) throws Exception { config.getAaConfig(), config.getExtraServerProperties(), config.getStoreVersionConfigOverride(), - config.getRecordTransformerFunction(), + config.getRecordTransformerConfig(), config.getOffsetRecord()); } @@ -827,7 +829,7 @@ private void runTest( AAConfig aaConfig, Map extraServerProperties, Consumer storeVersionConfigOverride, - DaVinciRecordTransformerFunctionalInterface recordTransformerFunction, + DaVinciRecordTransformerConfig recordTransformerConfig, OffsetRecord offsetRecord) throws Exception { int partitionCount = PARTITION_COUNT; @@ -856,7 +858,7 @@ private void runTest( diskUsageForTest, extraServerProperties, false, - recordTransformerFunction, + recordTransformerConfig, offsetRecord).build(); Properties kafkaProps = new Properties(); @@ -875,7 +877,7 @@ private void runTest( PARTITION_FOO, false, Optional.empty(), - recordTransformerFunction, + recordTransformerConfig, Lazy.of(() -> zkHelixAdmin))); Future testSubscribeTaskFuture = null; @@ -994,10 +996,10 @@ private StoreIngestionTaskFactory.Builder getIngestionTaskFactoryBuilder( Optional diskUsageForTest, Map extraServerProperties, Boolean isLiveConfigEnabled, - DaVinciRecordTransformerFunctionalInterface recordTransformerFunction, + DaVinciRecordTransformerConfig recordTransformerConfig, OffsetRecord optionalOffsetRecord) { - if (recordTransformerFunction != null) { + if (recordTransformerConfig != null && recordTransformerConfig.getRecordTransformerFunction() != null) { doReturn(mockAbstractStorageEngine).when(mockStorageEngineRepository).getLocalStorageEngine(topic); AbstractStorageIterator iterator = mock(AbstractStorageIterator.class); @@ -1007,8 +1009,9 @@ private StoreIngestionTaskFactory.Builder getIngestionTaskFactoryBuilder( when(mockAbstractStorageEngine.getIterator(anyInt())).thenReturn(iterator); } else { - doReturn(new DeepCopyStorageEngine(mockAbstractStorageEngine)).when(mockStorageEngineRepository) - .getLocalStorageEngine(topic); + mockDeepCopyStorageEngine = spy(new DeepCopyStorageEngine(mockAbstractStorageEngine)); + doReturn(mockDeepCopyStorageEngine).when(mockStorageEngineRepository).getLocalStorageEngine(topic); + doNothing().when(mockDeepCopyStorageEngine).createSnapshot(any()); } inMemoryLocalKafkaConsumer = @@ -2459,7 +2462,6 @@ public void testOffsetPersistent(AAConfig aaConfig) throws Exception { } finally { databaseSyncBytesIntervalForTransactionalMode = 1; } - } @Test(dataProvider = "aaConfigProvider") @@ -2497,6 +2499,48 @@ public void testVeniceMessagesProcessingWithSortedInput(AAConfig aaConfig) throw }, aaConfig); } + @Test(dataProvider = "True-and-False", dataProviderClass = DataProviderUtils.class) + public void testVeniceMessagesProcessingWithSortedInputWithBlobMode(boolean blobMode) throws Exception { + localVeniceWriter.broadcastStartOfPush(true, new HashMap<>()); + PubSubProduceResult putMetadata = + (PubSubProduceResult) localVeniceWriter.put(putKeyFoo, putValue, EXISTING_SCHEMA_ID).get(); + PubSubProduceResult deleteMetadata = (PubSubProduceResult) localVeniceWriter.delete(deleteKeyFoo, null).get(); + localVeniceWriter.broadcastEndOfPush(new HashMap<>()); + + StoreIngestionTaskTestConfig testConfig = new StoreIngestionTaskTestConfig(Utils.setOf(PARTITION_FOO), () -> { + // Verify it retrieves the offset from the Offset Manager + verify(mockStorageMetadataService, timeout(TEST_TIMEOUT_MS)).getLastOffset(topic, PARTITION_FOO); + + // Verify it commits the offset to Offset Manager after receiving EOP control message + OffsetRecord expectedOffsetRecordForDeleteMessage = getOffsetRecord(deleteMetadata.getOffset() + 1, true); + verify(mockStorageMetadataService, timeout(TEST_TIMEOUT_MS)) + .put(topic, PARTITION_FOO, expectedOffsetRecordForDeleteMessage); + // Deferred write is not going to commit offset for every message, but will commit offset for every control + // message + // The following verification is for START_OF_PUSH control message + verify(mockStorageMetadataService, times(1)) + .put(topic, PARTITION_FOO, getOffsetRecord(putMetadata.getOffset() - 1)); + // Check database mode switches from deferred-write to transactional after EOP control message + StoragePartitionConfig deferredWritePartitionConfig = new StoragePartitionConfig(topic, PARTITION_FOO); + deferredWritePartitionConfig.setDeferredWrite(!blobMode); + verify(mockAbstractStorageEngine, times(1)) + .beginBatchWrite(eq(deferredWritePartitionConfig), any(), eq(Optional.empty())); + StoragePartitionConfig transactionalPartitionConfig = new StoragePartitionConfig(topic, PARTITION_FOO); + verify(mockAbstractStorageEngine, times(1)).endBatchWrite(transactionalPartitionConfig); + }, null); + testConfig.setHybridStoreConfig( + Optional.of( + new HybridStoreConfigImpl( + 10, + 20, + HybridStoreConfigImpl.DEFAULT_HYBRID_TIME_LAG_THRESHOLD, + DataReplicationPolicy.NON_AGGREGATE, + BufferReplayPolicy.REWIND_FROM_EOP))); + testConfig.setExtraServerProperties( + Collections.singletonMap(RocksDBServerConfig.ROCKSDB_BLOB_FILES_ENABLED, Boolean.toString(blobMode))); + runTest(testConfig); + } + @Test(dataProvider = "aaConfigProvider") public void testVeniceMessagesProcessingWithSortedInputVerifyChecksum(AAConfig aaConfig) throws Exception { databaseChecksumVerificationEnabled = true; @@ -4769,15 +4813,16 @@ public void testStoreIngestionRecordTransformer(AAConfig aaConfig) throws Except when(leaderProducedRecordContext.getValueUnion()).thenReturn(put); when(leaderProducedRecordContext.getKeyBytes()).thenReturn(putKeyFoo); - Schema keySchema = Schema.create(Schema.Type.INT); + Schema myKeySchema = Schema.create(Schema.Type.INT); SchemaEntry keySchemaEntry = mock(SchemaEntry.class); - when(keySchemaEntry.getSchema()).thenReturn(keySchema); + when(keySchemaEntry.getSchema()).thenReturn(myKeySchema); when(mockSchemaRepo.getKeySchema(storeNameWithoutVersionInfo)).thenReturn(keySchemaEntry); - Schema valueSchema = Schema.create(Schema.Type.STRING); + Schema myValueSchema = Schema.create(Schema.Type.STRING); SchemaEntry valueSchemaEntry = mock(SchemaEntry.class); - when(valueSchemaEntry.getSchema()).thenReturn(valueSchema); + when(valueSchemaEntry.getSchema()).thenReturn(myValueSchema); when(mockSchemaRepo.getValueSchema(eq(storeNameWithoutVersionInfo), anyInt())).thenReturn(valueSchemaEntry); + when(mockSchemaRepo.getSupersetOrLatestValueSchema(eq(storeNameWithoutVersionInfo))).thenReturn(valueSchemaEntry); StoreIngestionTaskTestConfig config = new StoreIngestionTaskTestConfig(Collections.singleton(PARTITION_FOO), () -> { TestUtils.waitForNonDeterministicAssertion( @@ -4797,7 +4842,17 @@ public void testStoreIngestionRecordTransformer(AAConfig aaConfig) throws Except throw new VeniceException(e); } }, aaConfig); - config.setRecordTransformerFunction((storeVersion) -> new TestStringRecordTransformer(storeVersion, true)); + + DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> new TestStringRecordTransformer( + storeVersion, + keySchema, + inputValueSchema, + outputValueSchema, + true), + String.class, + Schema.create(Schema.Type.STRING)); + config.setRecordTransformerConfig(recordTransformerConfig); runTest(config); // Transformer error should never be recorded @@ -4831,15 +4886,16 @@ public void testStoreIngestionRecordTransformerError(AAConfig aaConfig) throws E when(leaderProducedRecordContext.getValueUnion()).thenReturn(put); when(leaderProducedRecordContext.getKeyBytes()).thenReturn(keyBytes); - Schema keySchema = Schema.create(Schema.Type.INT); + Schema myKeySchema = Schema.create(Schema.Type.INT); SchemaEntry keySchemaEntry = mock(SchemaEntry.class); - when(keySchemaEntry.getSchema()).thenReturn(keySchema); + when(keySchemaEntry.getSchema()).thenReturn(myKeySchema); when(mockSchemaRepo.getKeySchema(storeNameWithoutVersionInfo)).thenReturn(keySchemaEntry); - Schema valueSchema = Schema.create(Schema.Type.INT); + Schema myValueSchema = Schema.create(Schema.Type.INT); SchemaEntry valueSchemaEntry = mock(SchemaEntry.class); - when(valueSchemaEntry.getSchema()).thenReturn(valueSchema); + when(valueSchemaEntry.getSchema()).thenReturn(myValueSchema); when(mockSchemaRepo.getValueSchema(eq(storeNameWithoutVersionInfo), anyInt())).thenReturn(valueSchemaEntry); + when(mockSchemaRepo.getSupersetOrLatestValueSchema(eq(storeNameWithoutVersionInfo))).thenReturn(valueSchemaEntry); StoreIngestionTaskTestConfig config = new StoreIngestionTaskTestConfig(Collections.singleton(PARTITION_FOO), () -> { TestUtils.waitForNonDeterministicAssertion( @@ -4862,7 +4918,17 @@ public void testStoreIngestionRecordTransformerError(AAConfig aaConfig) throws E verify(mockVersionedStorageIngestionStats, timeout(1000)) .recordTransformerError(eq(storeNameWithoutVersionInfo), anyInt(), anyDouble(), anyLong()); }, aaConfig); - config.setRecordTransformerFunction((storeVersion) -> new TestStringRecordTransformer(storeVersion, true)); + + DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> new TestStringRecordTransformer( + storeVersion, + keySchema, + inputValueSchema, + outputValueSchema, + true), + String.class, + Schema.create(Schema.Type.STRING)); + config.setRecordTransformerConfig(recordTransformerConfig); runTest(config); } @@ -5181,27 +5247,24 @@ public void testMeasureLagWithCallToPubSub() { * {@link StoreIngestionTask#reportError(String, int, Exception)} should be called in order to trigger a Helix * state transition without waiting 24+ hours for the Helix state transition timeout. */ - @Test + @Test(timeOut = 30000) public void testProcessConsumerActionsError() throws Exception { runTest(Collections.singleton(PARTITION_FOO), () -> { // This is an actual exception thrown when deserializing a corrupted OffsetRecord String msg = "Received Magic Byte '6' which is not supported by InternalAvroSpecificSerializer. " + "The only supported Magic Byte for this implementation is '24'."; when(mockStorageMetadataService.getLastOffset(any(), anyInt())).thenThrow(new VeniceMessageException(msg)); - - for (int i = 0; i < StoreIngestionTask.MAX_CONSUMER_ACTION_ATTEMPTS; i++) { - try { - storeIngestionTaskUnderTest.processConsumerActions(storeAndVersionConfigsUnderTest.store); - } catch (InterruptedException e) { - throw new RuntimeException(e); - } + try { + storeIngestionTaskUnderTest.processConsumerActions(storeAndVersionConfigsUnderTest.store); + } catch (InterruptedException e) { + throw new RuntimeException(e); } + waitForNonDeterministicAssertion( + 30, + TimeUnit.SECONDS, + () -> assertTrue(storeIngestionTaskUnderTest.consumerActionsQueue.isEmpty(), "Wait until CAQ is empty")); ArgumentCaptor captor = ArgumentCaptor.forClass(VeniceException.class); - waitForNonDeterministicAssertion(15, TimeUnit.SECONDS, () -> { - assertTrue(storeIngestionTaskUnderTest.consumerActionsQueue.isEmpty(), "Wait for action processing"); - verify(storeIngestionTaskUnderTest, atLeastOnce()) - .reportError(anyString(), eq(PARTITION_FOO), captor.capture()); - }); + verify(storeIngestionTaskUnderTest, atLeastOnce()).reportError(anyString(), eq(PARTITION_FOO), captor.capture()); assertTrue(captor.getValue().getMessage().endsWith(msg)); }, AA_OFF); } @@ -5308,6 +5371,73 @@ public void testStoreIngestionInternalClose() throws Exception { dropPartitionFuture.get().get(); } + @Test(dataProviderClass = DataProviderUtils.class, dataProvider = "Two-True-and-False") + public void testSnapshotGenerationConditions(boolean isBlobTransferEnabled, boolean blobTransferManagerEnabled) { + Map serverProperties = new HashMap<>(); + serverProperties.put(BLOB_TRANSFER_MANAGER_ENABLED, blobTransferManagerEnabled); + + Version version = mock(Version.class); + doReturn(1).when(version).getPartitionCount(); + doReturn(VersionStatus.STARTED).when(version).getStatus(); + doReturn(true).when(version).isNativeReplicationEnabled(); + DataRecoveryVersionConfig dataRecoveryVersionConfig = new DataRecoveryVersionConfigImpl("dc-0", false, 1); + doReturn(dataRecoveryVersionConfig).when(version).getDataRecoveryVersionConfig(); + + StorageService storageService = mock(StorageService.class); + Store store = mock(Store.class); + + doReturn(version).when(store).getVersion(eq(1)); + + VeniceStoreVersionConfig storeConfig = mock(VeniceStoreVersionConfig.class); + doReturn(isBlobTransferEnabled).when(storeConfig).isBlobTransferEnabled(); + doReturn(topic).when(storeConfig).getStoreVersionName(); + + StoreIngestionTaskFactory ingestionTaskFactory = getIngestionTaskFactoryBuilder( + new RandomPollStrategy(), + Utils.setOf(PARTITION_FOO), + Optional.empty(), + serverProperties, + true, + null, + null).build(); + + doReturn(Version.parseStoreFromVersionTopic(topic)).when(store).getName(); + storeIngestionTaskUnderTest = ingestionTaskFactory.getNewIngestionTask( + storageService, + store, + version, + new Properties(), + isCurrentVersion, + storeConfig, + 1, + false, + Optional.empty(), + null, + null); + OffsetRecord offsetRecord = mock(OffsetRecord.class); + doReturn(pubSubTopic).when(offsetRecord).getLeaderTopic(any()); + doReturn(false).when(offsetRecord).isEndOfPushReceived(); + doReturn(100L).when(offsetRecord).getOffsetLag(); + PartitionConsumptionState partitionConsumptionState = + new PartitionConsumptionState(Utils.getReplicaId(pubSubTopic, 0), 0, offsetRecord, true); + + KafkaMessageEnvelope kafkaMessageEnvelope = spy(Mockito.mock(KafkaMessageEnvelope.class)); + ProducerMetadata producerMetadata = new ProducerMetadata(); + producerMetadata.producerGUID = GuidUtils.getGuidFromCharSequence("test_guid"); + producerMetadata.messageTimestamp = 1000L; + kafkaMessageEnvelope.producerMetadata = producerMetadata; + + // action + storeIngestionTaskUnderTest + .processEndOfPush(kafkaMessageEnvelope, 1, offsetRecord.getOffsetLag(), partitionConsumptionState); + // verify + if (isBlobTransferEnabled && blobTransferManagerEnabled) { + verify(mockDeepCopyStorageEngine).createSnapshot(any()); + } else { + verify(mockDeepCopyStorageEngine, never()).createSnapshot(any()); + } + } + private VeniceStoreVersionConfig getDefaultMockVeniceStoreVersionConfig( Consumer storeVersionConfigOverride) { // mock the store config diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/VeniceAdaptiveIngestionThrottlerTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/VeniceAdaptiveIngestionThrottlerTest.java new file mode 100644 index 00000000000..c9bc36cf035 --- /dev/null +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/kafka/consumer/VeniceAdaptiveIngestionThrottlerTest.java @@ -0,0 +1,35 @@ +package com.linkedin.davinci.kafka.consumer; + +import org.testng.Assert; +import org.testng.annotations.Test; + + +public class VeniceAdaptiveIngestionThrottlerTest { + @Test + public void testAdaptiveIngestionThrottler() { + VeniceAdaptiveIngestionThrottler adaptiveIngestionThrottler = + new VeniceAdaptiveIngestionThrottler(10, 100, 10, "test"); + adaptiveIngestionThrottler.registerLimiterSignal(() -> true); + adaptiveIngestionThrottler.checkSignalAndAdjustThrottler(); + Assert.assertEquals(adaptiveIngestionThrottler.getCurrentThrottlerIndex(), 2); + adaptiveIngestionThrottler.checkSignalAndAdjustThrottler(); + Assert.assertEquals(adaptiveIngestionThrottler.getCurrentThrottlerIndex(), 1); + adaptiveIngestionThrottler.checkSignalAndAdjustThrottler(); + Assert.assertEquals(adaptiveIngestionThrottler.getCurrentThrottlerIndex(), 0); + adaptiveIngestionThrottler.checkSignalAndAdjustThrottler(); + Assert.assertEquals(adaptiveIngestionThrottler.getCurrentThrottlerIndex(), 0); + adaptiveIngestionThrottler = new VeniceAdaptiveIngestionThrottler(10, 100, 10, "test"); + adaptiveIngestionThrottler.registerBoosterSignal(() -> true); + adaptiveIngestionThrottler.checkSignalAndAdjustThrottler(); + Assert.assertEquals(adaptiveIngestionThrottler.getCurrentThrottlerIndex(), 4); + + adaptiveIngestionThrottler = new VeniceAdaptiveIngestionThrottler(3, 100, 10, "test"); + + adaptiveIngestionThrottler.checkSignalAndAdjustThrottler(); + Assert.assertEquals(adaptiveIngestionThrottler.getCurrentThrottlerIndex(), 3); + adaptiveIngestionThrottler.checkSignalAndAdjustThrottler(); + Assert.assertEquals(adaptiveIngestionThrottler.getCurrentThrottlerIndex(), 3); + adaptiveIngestionThrottler.checkSignalAndAdjustThrottler(); + Assert.assertEquals(adaptiveIngestionThrottler.getCurrentThrottlerIndex(), 4); + } +} diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/repository/NativeMetadataRepositoryTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/repository/NativeMetadataRepositoryTest.java index fdf6b821a71..c9205c43ac4 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/repository/NativeMetadataRepositoryTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/repository/NativeMetadataRepositoryTest.java @@ -1,6 +1,7 @@ package com.linkedin.davinci.repository; import static com.linkedin.venice.ConfigKeys.CLIENT_SYSTEM_STORE_REPOSITORY_REFRESH_INTERVAL_SECONDS; +import static com.linkedin.venice.system.store.MetaStoreWriter.*; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doReturn; @@ -10,9 +11,12 @@ import com.linkedin.venice.client.store.ClientConfig; import com.linkedin.venice.client.store.schemas.TestKeyRecord; import com.linkedin.venice.client.store.schemas.TestValueRecord; +import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.exceptions.VeniceNoStoreException; import com.linkedin.venice.meta.Store; import com.linkedin.venice.meta.StoreConfig; +import com.linkedin.venice.schema.SchemaData; +import com.linkedin.venice.schema.SchemaEntry; import com.linkedin.venice.system.store.MetaStoreDataType; import com.linkedin.venice.systemstore.schemas.StoreKeySchemas; import com.linkedin.venice.systemstore.schemas.StoreMetaKey; @@ -24,6 +28,7 @@ import io.tehuti.Metric; import io.tehuti.metrics.MetricsRepository; import java.time.Clock; +import java.util.Collections; import java.util.HashMap; import java.util.Map; import java.util.concurrent.TimeUnit; @@ -190,6 +195,56 @@ protected Store fetchStoreFromRemote(String storeName, String clusterName) { } @Override + protected SchemaData getSchemaData(String storeName) { + SchemaData schemaData = schemaMap.get(storeName); + SchemaEntry keySchema; + if (schemaData == null) { + // Retrieve the key schema and initialize SchemaData only if it's not cached yet. + StoreMetaKey keySchemaKey = MetaStoreDataType.STORE_KEY_SCHEMAS + .getStoreMetaKey(Collections.singletonMap(KEY_STRING_STORE_NAME, storeName)); + Map keySchemaMap = + getStoreMetaValue(storeName, keySchemaKey).storeKeySchemas.keySchemaMap; + if (keySchemaMap.isEmpty()) { + throw new VeniceException("No key schema found for store: " + storeName); + } + Map.Entry keySchemaEntry = keySchemaMap.entrySet().iterator().next(); + keySchema = + new SchemaEntry(Integer.parseInt(keySchemaEntry.getKey().toString()), keySchemaEntry.getValue().toString()); + schemaData = new SchemaData(storeName, keySchema); + } + StoreMetaKey valueSchemaKey = MetaStoreDataType.STORE_VALUE_SCHEMAS + .getStoreMetaKey(Collections.singletonMap(KEY_STRING_STORE_NAME, storeName)); + Map valueSchemaMap = + getStoreMetaValue(storeName, valueSchemaKey).storeValueSchemas.valueSchemaMap; + // Check the value schema string, if it's empty then try to query the other key space for individual value schema. + for (Map.Entry entry: valueSchemaMap.entrySet()) { + // Check if we already have the corresponding value schema + int valueSchemaId = Integer.parseInt(entry.getKey().toString()); + if (schemaData.getValueSchema(valueSchemaId) != null) { + continue; + } + if (entry.getValue().toString().isEmpty()) { + // The value schemas might be too large to be stored in a single K/V. + StoreMetaKey individualValueSchemaKey = + MetaStoreDataType.STORE_VALUE_SCHEMA.getStoreMetaKey(new HashMap() { + { + put(KEY_STRING_STORE_NAME, storeName); + put(KEY_STRING_SCHEMA_ID, entry.getKey().toString()); + } + }); + // Empty string is not a valid value schema therefore it's safe to throw exceptions if we also cannot find it + // in + // the individual value schema key space. + String valueSchema = + getStoreMetaValue(storeName, individualValueSchemaKey).storeValueSchema.valueSchema.toString(); + schemaData.addValueSchema(new SchemaEntry(valueSchemaId, valueSchema)); + } else { + schemaData.addValueSchema(new SchemaEntry(valueSchemaId, entry.getValue().toString())); + } + } + return schemaData; + } + protected StoreMetaValue getStoreMetaValue(String storeName, StoreMetaKey key) { StoreMetaValue storeMetaValue = new StoreMetaValue(); MetaStoreDataType metaStoreDataType = MetaStoreDataType.valueOf(key.metadataType); diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/stats/DIVStatsReporterTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/stats/DIVStatsReporterTest.java index a5279bc6a9f..9c7a54254cb 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/stats/DIVStatsReporterTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/stats/DIVStatsReporterTest.java @@ -1,8 +1,9 @@ package com.linkedin.davinci.stats; import static com.linkedin.venice.stats.StatsErrorCode.NULL_DIV_STATS; -import static org.mockito.Mockito.*; -import static org.testng.Assert.*; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.testng.Assert.assertEquals; import com.linkedin.venice.tehuti.MockTehutiReporter; import com.linkedin.venice.utils.Utils; diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartitionTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartitionTest.java index f54c7dd885a..ba6e4e80f66 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartitionTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/store/rocksdb/RocksDBStoragePartitionTest.java @@ -1,11 +1,31 @@ package com.linkedin.davinci.store.rocksdb; -import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.*; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_BLOB_FILES_ENABLED; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_BLOB_FILE_SIZE_IN_BYTES; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_BLOB_FILE_STARTING_LEVEL; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_BLOCK_CACHE_IMPLEMENTATION; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_LEVEL0_COMPACTION_TUNING_FOR_READ_WRITE_LEADER_ENABLED; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_LEVEL0_FILE_NUM_COMPACTION_TRIGGER; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_LEVEL0_FILE_NUM_COMPACTION_TRIGGER_FOR_READ_WRITE_LEADER; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_LEVEL0_FILE_NUM_COMPACTION_TRIGGER_WRITE_ONLY_VERSION; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_LEVEL0_SLOWDOWN_WRITES_TRIGGER; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_LEVEL0_SLOWDOWN_WRITES_TRIGGER_FOR_READ_WRITE_LEADER; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_LEVEL0_SLOWDOWN_WRITES_TRIGGER_WRITE_ONLY_VERSION; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_LEVEL0_STOPS_WRITES_TRIGGER; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_LEVEL0_STOPS_WRITES_TRIGGER_FOR_READ_WRITE_LEADER; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_LEVEL0_STOPS_WRITES_TRIGGER_WRITE_ONLY_VERSION; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_MAX_MEMTABLE_COUNT; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_MEMTABLE_SIZE_IN_BYTES; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_MIN_BLOB_SIZE_IN_BYTES; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_PLAIN_TABLE_FORMAT_ENABLED; +import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_TOTAL_MEMTABLE_USAGE_CAP_IN_BYTES; +import static com.linkedin.venice.ConfigKeys.BLOB_TRANSFER_MANAGER_ENABLED; import static com.linkedin.venice.ConfigKeys.INGESTION_MEMORY_LIMIT; import static com.linkedin.venice.ConfigKeys.INGESTION_USE_DA_VINCI_CLIENT; import static com.linkedin.venice.ConfigKeys.PERSISTENCE_TYPE; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.mock; +import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; import static org.testng.AssertJUnit.assertFalse; @@ -22,6 +42,7 @@ import com.linkedin.venice.meta.PersistenceType; import com.linkedin.venice.meta.Version; import com.linkedin.venice.serialization.avro.AvroProtocolDefinition; +import com.linkedin.venice.store.rocksdb.RocksDBUtils; import com.linkedin.venice.utils.ByteUtils; import com.linkedin.venice.utils.DataProviderUtils; import com.linkedin.venice.utils.Utils; @@ -66,6 +87,10 @@ public class RocksDBStoragePartitionTest { BLOB_GARBAGE_METRIC); private Map generateInput(int recordCnt, boolean sorted, int padLength) { + return generateInput(recordCnt, sorted, padLength, 0); + } + + private Map generateInput(int recordCnt, boolean sorted, int padLength, int startIdx) { Map records; if (sorted) { BytewiseComparator comparator = new BytewiseComparator(new ComparatorOptions()); @@ -77,7 +102,7 @@ private Map generateInput(int recordCnt, boolean sorted, int pad } else { records = new HashMap<>(); } - for (int i = 0; i < recordCnt; ++i) { + for (int i = startIdx; i < recordCnt + startIdx; ++i) { String value = VALUE_PREFIX + i; if (padLength > 0) { value += RandomStringUtils.random(padLength, true, true); @@ -118,6 +143,131 @@ public Object[][] testIngestionDataProvider() { }; } + @Test + public void testBlobDBCompatibility() { + String storeName = Version.composeKafkaTopic(Utils.getUniqueString("test_store"), 1); + String storeDir = getTempDatabaseDir(storeName); + int partitionId = 0; + String dbFolder = RocksDBUtils.composePartitionDbDir(DATA_BASE_DIR, storeName, partitionId); + File dbDir = new File(dbFolder); + + Supplier sstFileFinder = () -> dbDir.list(((dir, name) -> name.endsWith(".sst"))); + Supplier blobFileFinder = () -> dbDir.list(((dir, name) -> name.endsWith(".blob"))); + + StoragePartitionConfig partitionConfig = new StoragePartitionConfig(storeName, partitionId); + + Map largeInputRecords = generateInput(1000, false, 10000, 0); + Map smallInputRecords = generateInput(1000, false, 10, 10000); + List> largeEntryList = new ArrayList<>(largeInputRecords.entrySet()); + List> smallEntryList = new ArrayList<>(smallInputRecords.entrySet()); + Properties extraProps = new Properties(); + // Disable blob files + extraProps.put(ROCKSDB_BLOB_FILES_ENABLED, "false"); + extraProps.put(ROCKSDB_MIN_BLOB_SIZE_IN_BYTES, "1000"); // make sure the threshold is larger than small records + // generated + extraProps.put(ROCKSDB_BLOB_FILE_SIZE_IN_BYTES, "2097152"); + extraProps.put(ROCKSDB_BLOB_FILE_STARTING_LEVEL, "0"); + extraProps.put(ROCKSDB_MEMTABLE_SIZE_IN_BYTES, "1048576"); // 1MB + + VeniceProperties veniceServerProperties = + AbstractStorageEngineTest.getServerProperties(PersistenceType.ROCKS_DB, extraProps); + RocksDBServerConfig rocksDBServerConfig = new RocksDBServerConfig(veniceServerProperties); + VeniceServerConfig serverConfig = new VeniceServerConfig(veniceServerProperties); + RocksDBStorageEngineFactory factory = new RocksDBStorageEngineFactory(serverConfig); + VeniceStoreVersionConfig storeConfig = new VeniceStoreVersionConfig(storeName, veniceServerProperties); + RocksDBStoragePartition storagePartition = new RocksDBStoragePartition( + partitionConfig, + factory, + DATA_BASE_DIR, + null, + ROCKSDB_THROTTLER, + rocksDBServerConfig, + storeConfig); + // Insert the first 300 [0, 300) entries with blob db disabled + for (int i = 0; i < 300; i++) { + storagePartition.put(largeEntryList.get(i).getKey().getBytes(), largeEntryList.get(i).getValue().getBytes()); + storagePartition.put(smallEntryList.get(i).getKey().getBytes(), smallEntryList.get(i).getValue().getBytes()); + } + storagePartition.close(); + // Make sure no blob files were generated + assertTrue(sstFileFinder.get().length > 0); + assertTrue(blobFileFinder.get().length == 0); + + // Enable blob files + extraProps.put(ROCKSDB_BLOB_FILES_ENABLED, "true"); + + veniceServerProperties = AbstractStorageEngineTest.getServerProperties(PersistenceType.ROCKS_DB, extraProps); + rocksDBServerConfig = new RocksDBServerConfig(veniceServerProperties); + serverConfig = new VeniceServerConfig(veniceServerProperties); + factory = new RocksDBStorageEngineFactory(serverConfig); + storeConfig = new VeniceStoreVersionConfig(storeName, veniceServerProperties); + storagePartition = new RocksDBStoragePartition( + partitionConfig, + factory, + DATA_BASE_DIR, + null, + ROCKSDB_THROTTLER, + rocksDBServerConfig, + storeConfig); + // Insert [300, 700) entries with blob db enabled + for (int i = 300; i < 700; i++) { + storagePartition.put(largeEntryList.get(i).getKey().getBytes(), largeEntryList.get(i).getValue().getBytes()); + storagePartition.put(smallEntryList.get(i).getKey().getBytes(), smallEntryList.get(i).getValue().getBytes()); + } + storagePartition.sync(); + // Make sure blob files were generated + assertTrue(sstFileFinder.get().length > 0); + int blobFileCnt = blobFileFinder.get().length; + assertTrue(blobFileCnt > 0); + // Validate all the entries inserted so far + for (int i = 0; i < 700; i++) { + Assert.assertEquals( + storagePartition.get(largeEntryList.get(i).getKey().getBytes()), + largeEntryList.get(i).getValue().getBytes()); + Assert.assertEquals( + storagePartition.get(smallEntryList.get(i).getKey().getBytes()), + smallEntryList.get(i).getValue().getBytes()); + } + storagePartition.close(); + + // Disable blob files + extraProps.put(ROCKSDB_BLOB_FILES_ENABLED, "false"); + + veniceServerProperties = AbstractStorageEngineTest.getServerProperties(PersistenceType.ROCKS_DB, extraProps); + rocksDBServerConfig = new RocksDBServerConfig(veniceServerProperties); + serverConfig = new VeniceServerConfig(veniceServerProperties); + factory = new RocksDBStorageEngineFactory(serverConfig); + storeConfig = new VeniceStoreVersionConfig(storeName, veniceServerProperties); + storagePartition = new RocksDBStoragePartition( + partitionConfig, + factory, + DATA_BASE_DIR, + null, + ROCKSDB_THROTTLER, + rocksDBServerConfig, + storeConfig); + // Insert [700, 1000) entries with blob db enabled + for (int i = 700; i < 1000; i++) { + storagePartition.put(largeEntryList.get(i).getKey().getBytes(), largeEntryList.get(i).getValue().getBytes()); + storagePartition.put(smallEntryList.get(i).getKey().getBytes(), smallEntryList.get(i).getValue().getBytes()); + } + + storagePartition.sync(); + // Make sure no new blob files were generated + assertEquals(blobFileFinder.get().length, blobFileCnt); + // Validate all the entries inserted previously + for (Map.Entry entry: largeEntryList) { + Assert.assertEquals(storagePartition.get(entry.getKey().getBytes()), entry.getValue().getBytes()); + } + for (Map.Entry entry: smallEntryList) { + Assert.assertEquals(storagePartition.get(entry.getKey().getBytes()), entry.getValue().getBytes()); + } + + storagePartition.close(); + storagePartition.drop(); + removeDir(storeDir); + } + @Test(dataProvider = "testIngestionDataProvider") public void testIngestion( boolean sorted, @@ -962,7 +1112,10 @@ public void testCreateSnapshot(boolean blobTransferEnabled) { int partitionId = 0; StoragePartitionConfig partitionConfig = new StoragePartitionConfig(storeName, partitionId); partitionConfig.setDeferredWrite(false); - VeniceProperties veniceServerProperties = AbstractStorageEngineTest.getServerProperties(PersistenceType.ROCKS_DB); + Properties extraProps = new Properties(); + extraProps.setProperty(BLOB_TRANSFER_MANAGER_ENABLED, "true"); + VeniceProperties veniceServerProperties = + AbstractStorageEngineTest.getServerProperties(PersistenceType.ROCKS_DB, extraProps); RocksDBServerConfig rocksDBServerConfig = new RocksDBServerConfig(veniceServerProperties); VeniceServerConfig serverConfig = new VeniceServerConfig(veniceServerProperties); diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/transformer/RecordTransformerTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/transformer/RecordTransformerTest.java index 1cb8ef654ae..0645f2be9ba 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/transformer/RecordTransformerTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/transformer/RecordTransformerTest.java @@ -40,15 +40,15 @@ public void deleteClassHash() { @Test public void testRecordTransformer() { + Schema keySchema = Schema.create(Schema.Type.INT); + Schema valueSchema = Schema.create(Schema.Type.STRING); + DaVinciRecordTransformer recordTransformer = - new TestStringRecordTransformer(storeVersion, false); + new TestStringRecordTransformer(storeVersion, keySchema, valueSchema, valueSchema, false); assertEquals(recordTransformer.getStoreVersion(), storeVersion); - Schema keySchema = recordTransformer.getKeySchema(); - assertEquals(keySchema.getType(), Schema.Type.INT); - - Schema outputValueSchema = recordTransformer.getOutputValueSchema(); - assertEquals(outputValueSchema.getType(), Schema.Type.STRING); + assertEquals(recordTransformer.getKeySchema().getType(), Schema.Type.INT); + assertEquals(recordTransformer.getOutputValueSchema().getType(), Schema.Type.STRING); Lazy lazyKey = Lazy.of(() -> 42); Lazy lazyValue = Lazy.of(() -> "SampleValue"); @@ -72,8 +72,12 @@ public void testRecordTransformer() { @Test public void testOnRecovery() { + Schema keySchema = Schema.create(Schema.Type.INT); + Schema valueSchema = Schema.create(Schema.Type.STRING); + DaVinciRecordTransformer recordTransformer = - new TestStringRecordTransformer(storeVersion, true); + new TestStringRecordTransformer(storeVersion, keySchema, valueSchema, valueSchema, true); + assertEquals(recordTransformer.getStoreVersion(), storeVersion); AbstractStorageIterator iterator = mock(AbstractStorageIterator.class); when(iterator.isValid()).thenReturn(true).thenReturn(false); @@ -99,18 +103,26 @@ public void testOnRecovery() { @Test public void testBlockingRecordTransformer() { - DaVinciRecordTransformer recordTransformer = new TestStringRecordTransformer(0, true); - recordTransformer = - new BlockingDaVinciRecordTransformer<>(recordTransformer, recordTransformer.getStoreRecordsInDaVinci()); - recordTransformer.onStartVersionIngestion(); + Schema keySchema = Schema.create(Schema.Type.INT); + Schema valueSchema = Schema.create(Schema.Type.STRING); + + DaVinciRecordTransformer recordTransformer = + new TestStringRecordTransformer(storeVersion, keySchema, valueSchema, valueSchema, true); + assertEquals(recordTransformer.getStoreVersion(), storeVersion); + + recordTransformer = new BlockingDaVinciRecordTransformer<>( + recordTransformer, + keySchema, + valueSchema, + valueSchema, + recordTransformer.getStoreRecordsInDaVinci()); + recordTransformer.onStartVersionIngestion(true); assertTrue(recordTransformer.getStoreRecordsInDaVinci()); - Schema keySchema = recordTransformer.getKeySchema(); - assertEquals(keySchema.getType(), Schema.Type.INT); + assertEquals(recordTransformer.getKeySchema().getType(), Schema.Type.INT); - Schema outputValueSchema = recordTransformer.getOutputValueSchema(); - assertEquals(outputValueSchema.getType(), Schema.Type.STRING); + assertEquals(recordTransformer.getOutputValueSchema().getType(), Schema.Type.STRING); Lazy lazyKey = Lazy.of(() -> 42); Lazy lazyValue = Lazy.of(() -> "SampleValue"); @@ -120,7 +132,7 @@ public void testBlockingRecordTransformer() { recordTransformer.processDelete(lazyKey); - recordTransformer.onEndVersionIngestion(); + recordTransformer.onEndVersionIngestion(2); } } diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/transformer/TestStringRecordTransformer.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/transformer/TestStringRecordTransformer.java index f9d7a3020b5..c8076da06de 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/transformer/TestStringRecordTransformer.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/transformer/TestStringRecordTransformer.java @@ -8,18 +8,13 @@ public class TestStringRecordTransformer extends DaVinciRecordTransformer { - public TestStringRecordTransformer(int storeVersion, boolean storeRecordsInDaVinci) { - super(storeVersion, storeRecordsInDaVinci); - } - - @Override - public Schema getKeySchema() { - return Schema.create(Schema.Type.INT); - } - - @Override - public Schema getOutputValueSchema() { - return Schema.create(Schema.Type.STRING); + public TestStringRecordTransformer( + int storeVersion, + Schema keySchema, + Schema inputValueSchema, + Schema outputValueSchema, + boolean storeRecordsInDaVinci) { + super(storeVersion, keySchema, inputValueSchema, outputValueSchema, storeRecordsInDaVinci); } @Override diff --git a/clients/venice-admin-tool/build.gradle b/clients/venice-admin-tool/build.gradle index d24d7a84a59..c73ea9b2819 100644 --- a/clients/venice-admin-tool/build.gradle +++ b/clients/venice-admin-tool/build.gradle @@ -44,6 +44,11 @@ jar { } } +shadowJar { + // Enable merging service files from different dependencies. Required to make gRPC based clients work. + mergeServiceFiles() +} + ext { jacocoCoverageThreshold = 0.00 } diff --git a/clients/venice-admin-tool/src/main/java/com/linkedin/venice/AdminTool.java b/clients/venice-admin-tool/src/main/java/com/linkedin/venice/AdminTool.java index f9467b2ce9b..b45179c581a 100644 --- a/clients/venice-admin-tool/src/main/java/com/linkedin/venice/AdminTool.java +++ b/clients/venice-admin-tool/src/main/java/com/linkedin/venice/AdminTool.java @@ -117,6 +117,7 @@ import com.linkedin.venice.utils.Time; import com.linkedin.venice.utils.Utils; import com.linkedin.venice.utils.VeniceProperties; +import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap; import com.linkedin.venice.utils.pools.LandFillObjectPool; import java.io.BufferedReader; import java.io.Console; @@ -125,6 +126,7 @@ import java.io.InputStreamReader; import java.nio.charset.StandardCharsets; import java.nio.file.Files; +import java.nio.file.Path; import java.nio.file.Paths; import java.time.Duration; import java.time.LocalDateTime; @@ -143,6 +145,10 @@ import java.util.Properties; import java.util.Set; import java.util.StringJoiner; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; import java.util.concurrent.TimeUnit; import java.util.function.Consumer; import java.util.function.Function; @@ -588,6 +594,9 @@ public static void main(String[] args) throws Exception { case DUMP_HOST_HEARTBEAT: dumpHostHeartbeat(cmd); break; + case CLUSTER_BATCH_TASK: + clusterBatchTask(cmd); + break; default: StringJoiner availableCommands = new StringJoiner(", "); for (Command c: Command.values()) { @@ -852,6 +861,96 @@ private static void deleteStore(CommandLine cmd) throws IOException { printObject(response); } + private static void clusterBatchTask(CommandLine cmd) { + String clusterName = getRequiredArgument(cmd, Arg.CLUSTER, Command.CLUSTER_BATCH_TASK); + String task = getRequiredArgument(cmd, Arg.TASK_NAME, Command.CLUSTER_BATCH_TASK); + String checkpointFile = getRequiredArgument(cmd, Arg.CHECKPOINT_FILE, Command.CLUSTER_BATCH_TASK); + int parallelism = Integer.parseInt(getOptionalArgument(cmd, Arg.THREAD_COUNT, "1")); + LOGGER.info( + "[**** Cluster Command Params ****] Cluster: {}, Task: {}, Checkpoint: {}, Parallelism: {}", + clusterName, + task, + checkpointFile, + parallelism); + // Create child data center controller client map. + ChildAwareResponse childAwareResponse = controllerClient.listChildControllers(clusterName); + Map controllerClientMap = getControllerClientMap(clusterName, childAwareResponse); + + // Fetch list cluster store list from parent region. + Map progressMap = new VeniceConcurrentHashMap<>(); + MultiStoreResponse clusterStoreResponse = controllerClient.queryStoreList(false); + if (clusterStoreResponse.isError()) { + throw new VeniceException("Unable to fetch cluster store list: " + clusterStoreResponse.getError()); + } + for (String storeName: clusterStoreResponse.getStores()) { + progressMap.put(storeName, Boolean.FALSE); + } + + // Load progress from checkpoint file. If file does not exist, it will create new one during checkpointing. + try { + Path checkpointFilePath = Paths.get(checkpointFile); + if (!Files.exists(checkpointFilePath.toAbsolutePath())) { + LOGGER.info( + "Checkpoint file path does not exist, will create a new checkpoint file: {}", + checkpointFilePath.toAbsolutePath()); + } else { + List fileLines = Files.readAllLines(checkpointFilePath); + for (String line: fileLines) { + String storeName = line.split(",")[0]; + // For now, it is boolean to start with, we can add more states to support retry. + boolean status = false; + if (line.split(",").length > 1) { + status = Boolean.parseBoolean(line.split(",")[1]); + } + progressMap.put(storeName, status); + } + } + } catch (IOException e) { + throw new VeniceException(e); + } + List taskList = + progressMap.entrySet().stream().filter(e -> !e.getValue()).map(Map.Entry::getKey).collect(Collectors.toList()); + + // Validate task type. For now, we only has one task, if we have more task in the future, we can extend this logic. + Supplier> functionSupplier = null; + if (SystemStorePushTask.TASK_NAME.equals(task)) { + String systemStoreType = getOptionalArgument(cmd, Arg.SYSTEM_STORE_TYPE); + if (systemStoreType != null) { + if (!(systemStoreType.equalsIgnoreCase(VeniceSystemStoreType.DAVINCI_PUSH_STATUS_STORE.toString()) + || systemStoreType.equalsIgnoreCase(VeniceSystemStoreType.META_STORE.toString()))) { + printErrAndExit("System store type: " + systemStoreType + " is not supported."); + } + } + System.out.println( + functionSupplier = () -> new SystemStorePushTask( + controllerClient, + controllerClientMap, + clusterName, + systemStoreType == null ? Optional.empty() : Optional.of(systemStoreType))); + } else { + printErrAndExit("Undefined task: " + task); + } + + // Create thread pool and start parallel processing. + ExecutorService executorService = Executors.newFixedThreadPool(parallelism); + List futureList = new ArrayList<>(); + for (int i = 0; i < parallelism; i++) { + BatchMaintenanceTaskRunner batchMaintenanceTaskRunner = + new BatchMaintenanceTaskRunner(progressMap, checkpointFile, taskList, functionSupplier.get()); + futureList.add(executorService.submit(batchMaintenanceTaskRunner)); + } + for (int i = 0; i < parallelism; i++) { + try { + futureList.get(i).get(); + LOGGER.info("Cluster task completed for thread : {}", i); + } catch (InterruptedException | ExecutionException e) { + LOGGER.warn(e.getMessage()); + executorService.shutdownNow(); + } + } + executorService.shutdownNow(); + } + private static void backfillSystemStores(CommandLine cmd) { String clusterName = getRequiredArgument(cmd, Arg.CLUSTER, Command.BACKFILL_SYSTEM_STORES); String systemStoreType = getRequiredArgument(cmd, Arg.SYSTEM_STORE_TYPE, Command.BACKFILL_SYSTEM_STORES); diff --git a/clients/venice-admin-tool/src/main/java/com/linkedin/venice/Arg.java b/clients/venice-admin-tool/src/main/java/com/linkedin/venice/Arg.java index 75b2115a495..05603f884c4 100644 --- a/clients/venice-admin-tool/src/main/java/com/linkedin/venice/Arg.java +++ b/clients/venice-admin-tool/src/main/java/com/linkedin/venice/Arg.java @@ -226,7 +226,10 @@ public enum Arg { SYSTEM_STORE_TYPE( "system-store-type", "sst", true, "Type of system store to backfill. Supported types are davinci_push_status_store and meta_store" - ), RETRY("retry", "r", false, "Retry this operation"), + ), TASK_NAME("task-name", "tn", true, "Name of the task for cluster command. Supported command [PushSystemStore]."), + CHECKPOINT_FILE("checkpoint-file", "cf", true, "Checkpoint file path for cluster command."), + THREAD_COUNT("thread-count", "tc", true, "Number of threads to execute. 1 if not specified"), + RETRY("retry", "r", false, "Retry this operation"), DISABLE_LOG("disable-log", "dl", false, "Disable logs from internal classes. Only print command output on console"), STORE_VIEW_CONFIGS( "storage-view-configs", "svc", true, diff --git a/clients/venice-admin-tool/src/main/java/com/linkedin/venice/BatchMaintenanceTaskRunner.java b/clients/venice-admin-tool/src/main/java/com/linkedin/venice/BatchMaintenanceTaskRunner.java new file mode 100644 index 00000000000..c8241442b01 --- /dev/null +++ b/clients/venice-admin-tool/src/main/java/com/linkedin/venice/BatchMaintenanceTaskRunner.java @@ -0,0 +1,100 @@ +package com.linkedin.venice; + +import com.linkedin.venice.exceptions.VeniceException; +import java.io.IOException; +import java.nio.file.Files; +import java.nio.file.Paths; +import java.util.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Function; +import java.util.stream.Collectors; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * This class is a simple runnable which keeps fetching task from list and execute the assigned task. The task fetching + * and progress tracking / checkpointing is thread-safe, so it can be run in parallel. + */ +public class BatchMaintenanceTaskRunner implements Runnable { + private static final Logger LOGGER = LogManager.getLogger(BatchMaintenanceTaskRunner.class); + private static final String TASK_LOG_PREFIX = "[**** TASK INFO ****]"; + + private static final ReentrantLock LOCK = new ReentrantLock(); + private static final AtomicInteger INDEX = new AtomicInteger(-1); + private final List taskList; + private final Function storeRunnable; + private final Map progressMap; + private final String checkpointFile; + + public BatchMaintenanceTaskRunner( + Map progressMap, + String checkpointFile, + List taskList, + Function storeRunnable) { + this.taskList = taskList; + this.storeRunnable = storeRunnable; + this.progressMap = progressMap; + this.checkpointFile = checkpointFile; + } + + @Override + public void run() { + while (true) { + int fetchedTaskIndex = INDEX.incrementAndGet(); + if (fetchedTaskIndex >= taskList.size()) { + LOGGER.info("Cannot find new store from queue, will exit."); + break; + } + String store = taskList.get(fetchedTaskIndex); + try { + LOGGER.info("{} Running store job: {} for store: {}", TASK_LOG_PREFIX, fetchedTaskIndex + 1, store); + boolean result = storeRunnable.apply(store); + if (result) { + LOGGER.info( + "{} Complete store task for job: {}/{} store: {}", + TASK_LOG_PREFIX, + fetchedTaskIndex + 1, + taskList.size(), + store); + progressMap.put(store, true); + } else { + LOGGER.info( + "{} Failed store task for job: {}/{} store: {}", + TASK_LOG_PREFIX, + fetchedTaskIndex + 1, + taskList.size(), + store); + } + // Periodically update the checkpoint file. + if ((fetchedTaskIndex % 100) == 0) { + LOGGER.info("{} Preparing to checkpoint status at index {}", TASK_LOG_PREFIX, fetchedTaskIndex); + checkpoint(checkpointFile); + } + } catch (Exception e) { + LOGGER.info("{} Caught exception: {}. Will exit.", TASK_LOG_PREFIX, e.getMessage()); + } + } + // Perform one final checkpointing before existing the runnable. + checkpoint(checkpointFile); + } + + public void checkpoint(String checkpointFile) { + try { + LOCK.lock(); + LOGGER.info("Updating checkpoint..."); + + List status = + progressMap.entrySet().stream().map(e -> e.getKey() + "," + e.getValue()).collect(Collectors.toList()); + Files.write(Paths.get(checkpointFile), status); + LOGGER.info("Updated checkpoint..."); + + } catch (IOException e) { + throw new VeniceException(e); + } finally { + LOCK.unlock(); + } + } +} diff --git a/clients/venice-admin-tool/src/main/java/com/linkedin/venice/Command.java b/clients/venice-admin-tool/src/main/java/com/linkedin/venice/Command.java index a943c76a3ed..abdf6089278 100644 --- a/clients/venice-admin-tool/src/main/java/com/linkedin/venice/Command.java +++ b/clients/venice-admin-tool/src/main/java/com/linkedin/venice/Command.java @@ -12,6 +12,7 @@ import static com.linkedin.venice.Arg.BATCH_GET_LIMIT; import static com.linkedin.venice.Arg.BLOB_TRANSFER_ENABLED; import static com.linkedin.venice.Arg.BOOTSTRAP_TO_ONLINE_TIMEOUT_IN_HOUR; +import static com.linkedin.venice.Arg.CHECKPOINT_FILE; import static com.linkedin.venice.Arg.CHILD_CONTROLLER_ADMIN_TOPIC_CONSUMPTION_ENABLED; import static com.linkedin.venice.Arg.CHUNKING_ENABLED; import static com.linkedin.venice.Arg.CLIENT_DECOMPRESSION_ENABLED; @@ -127,6 +128,8 @@ import static com.linkedin.venice.Arg.SYSTEM_STORE_TYPE; import static com.linkedin.venice.Arg.TARGET_SWAP_REGION; import static com.linkedin.venice.Arg.TARGET_SWAP_REGION_WAIT_TIME; +import static com.linkedin.venice.Arg.TASK_NAME; +import static com.linkedin.venice.Arg.THREAD_COUNT; import static com.linkedin.venice.Arg.TO_BE_STOPPED_NODES; import static com.linkedin.venice.Arg.UNUSED_SCHEMA_DELETION_ENABLED; import static com.linkedin.venice.Arg.URL; @@ -209,6 +212,10 @@ public enum Command { "backfill-system-stores", "Create system stores of a given type for user stores in a cluster", new Arg[] { URL, CLUSTER, SYSTEM_STORE_TYPE } ), + CLUSTER_BATCH_TASK( + "cluster-batch-task", "Run specific task against all user stores in a cluster in parallel", + new Arg[] { URL, CLUSTER, TASK_NAME, CHECKPOINT_FILE }, new Arg[] { THREAD_COUNT } + ), SET_VERSION( "set-version", "Set the version that will be served", new Arg[] { URL, STORE, VERSION }, new Arg[] { CLUSTER } ), ADD_SCHEMA("add-schema", "", new Arg[] { URL, STORE, VALUE_SCHEMA }, new Arg[] { CLUSTER }), diff --git a/clients/venice-admin-tool/src/main/java/com/linkedin/venice/SystemStorePushTask.java b/clients/venice-admin-tool/src/main/java/com/linkedin/venice/SystemStorePushTask.java new file mode 100644 index 00000000000..b214c46afc4 --- /dev/null +++ b/clients/venice-admin-tool/src/main/java/com/linkedin/venice/SystemStorePushTask.java @@ -0,0 +1,175 @@ +package com.linkedin.venice; + +import static com.linkedin.venice.AdminTool.printObject; + +import com.linkedin.venice.common.VeniceSystemStoreType; +import com.linkedin.venice.controllerapi.ControllerClient; +import com.linkedin.venice.controllerapi.ControllerResponse; +import com.linkedin.venice.controllerapi.JobStatusQueryResponse; +import com.linkedin.venice.controllerapi.StoreResponse; +import com.linkedin.venice.controllerapi.UpdateStoreQueryParams; +import com.linkedin.venice.controllerapi.VersionCreationResponse; +import com.linkedin.venice.controllerapi.VersionResponse; +import com.linkedin.venice.meta.Version; +import com.linkedin.venice.pushmonitor.ExecutionStatus; +import com.linkedin.venice.utils.Utils; +import java.util.Arrays; +import java.util.Collections; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.TimeUnit; +import java.util.function.Function; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * This class aims to do one time emtpy push to all user system stores of a specific user store. + * It will aggregate and compute the largest used version from all regions and update store before performing empty push. + * It will also skip empty push to store which is being migrated and is in the destination cluster. + */ +public class SystemStorePushTask implements Function { + public static final String TASK_NAME = "PushSystemStore"; + private static final Logger LOGGER = LogManager.getLogger(SystemStorePushTask.class); + private static final int JOB_POLLING_RETRY_COUNT = 200; + private static final int JOB_POLLING_RETRY_PERIOD_IN_SECONDS = 5; + private static final String SYSTEM_STORE_PUSH_TASK_LOG_PREFIX = "[**** SYSTEM STORE PUSH ****]"; + private final List systemStoreTypeList; + + private final ControllerClient parentControllerClient; + private final String clusterName; + private final Map childControllerClientMap; + + public SystemStorePushTask( + ControllerClient parentControllerClient, + Map controllerClientMap, + String clusterName, + Optional systemStoreTypeFilter) { + this.parentControllerClient = parentControllerClient; + this.childControllerClientMap = controllerClientMap; + this.clusterName = clusterName; + if (systemStoreTypeFilter.isPresent()) { + systemStoreTypeList = Collections.singletonList(VeniceSystemStoreType.valueOf(systemStoreTypeFilter.get())); + } else { + systemStoreTypeList = + Arrays.asList(VeniceSystemStoreType.META_STORE, VeniceSystemStoreType.DAVINCI_PUSH_STATUS_STORE); + } + } + + public Boolean apply(String storeName) { + StoreResponse storeResponse = parentControllerClient.getStore(storeName); + if (storeResponse.isError()) { + LOGGER.error("{} Unable to locate user store: {}", SYSTEM_STORE_PUSH_TASK_LOG_PREFIX, storeName); + return false; + } + if (storeResponse.getStore().isMigrating() && storeResponse.getStore().isMigrationDuplicateStore()) { + LOGGER.error( + "{} Unable to empty push to system store of migrating dest cluster store: {} in cluster: {}", + SYSTEM_STORE_PUSH_TASK_LOG_PREFIX, + storeName, + clusterName); + return false; + } + + for (VeniceSystemStoreType type: systemStoreTypeList) { + String systemStoreName = type.getSystemStoreName(storeName); + /** + * In current implementation, a push to system store will flip the flag to true, which can introduce unexpected + * behavior to the store. Here, we skip the system store push if it is turned off. + */ + boolean isSystemStoreEnabled = VeniceSystemStoreType.META_STORE.equals(type) + ? storeResponse.getStore().isStoreMetaSystemStoreEnabled() + : storeResponse.getStore().isDaVinciPushStatusStoreEnabled(); + if (!isSystemStoreEnabled) { + LOGGER.warn( + "{} System store: {} is disabled. Will skip the push.", + SYSTEM_STORE_PUSH_TASK_LOG_PREFIX, + systemStoreName); + } + VersionResponse response = parentControllerClient.getStoreLargestUsedVersion(clusterName, systemStoreName); + if (response.isError()) { + LOGGER.error( + "{} Unable to locate largest used store version for: {}", + SYSTEM_STORE_PUSH_TASK_LOG_PREFIX, + systemStoreName); + return false; + } + int largestUsedVersion = response.getVersion(); + + int version = getStoreLargestUsedVersionNumber(parentControllerClient, systemStoreName); + if (version == -1) { + return false; + } + largestUsedVersion = Math.max(largestUsedVersion, version); + for (Map.Entry controllerClientEntry: childControllerClientMap.entrySet()) { + int result = getStoreLargestUsedVersionNumber(controllerClientEntry.getValue(), systemStoreName); + if (result == -1) { + LOGGER.error( + "{} Unable to locate store for: {} in region: {}", + SYSTEM_STORE_PUSH_TASK_LOG_PREFIX, + systemStoreName, + controllerClientEntry.getKey()); + return false; + } + largestUsedVersion = Math.max(largestUsedVersion, result); + } + + LOGGER.info("Aggregate largest version: {} for store: {}", largestUsedVersion, systemStoreName); + ControllerResponse controllerResponse = parentControllerClient + .updateStore(systemStoreName, new UpdateStoreQueryParams().setLargestUsedVersionNumber(largestUsedVersion)); + if (controllerResponse.isError()) { + LOGGER.error( + "{} Unable to set largest used store version for: {} as {} in all regions", + SYSTEM_STORE_PUSH_TASK_LOG_PREFIX, + systemStoreName, + largestUsedVersion); + return false; + } + + VersionCreationResponse versionCreationResponse = + parentControllerClient.emptyPush(systemStoreName, "SYSTEM_STORE_PUSH_" + System.currentTimeMillis(), 10000); + // Kafka topic name in the above response is null, and it will be fixed with this code change. + String topicName = Version.composeKafkaTopic(systemStoreName, versionCreationResponse.getVersion()); + // Polling job status to make sure the empty push hits every child colo + int count = JOB_POLLING_RETRY_COUNT; + while (true) { + JobStatusQueryResponse jobStatusQueryResponse = + parentControllerClient.retryableRequest(3, controllerClient -> controllerClient.queryJobStatus(topicName)); + printObject(jobStatusQueryResponse, System.out::print); + if (jobStatusQueryResponse.isError()) { + return false; + } + ExecutionStatus executionStatus = ExecutionStatus.valueOf(jobStatusQueryResponse.getStatus()); + if (executionStatus.isTerminal()) { + if (executionStatus.isError()) { + LOGGER.error("{} Push error for topic: {}", SYSTEM_STORE_PUSH_TASK_LOG_PREFIX, topicName); + return false; + } + LOGGER.info("{} Push completed: {}", SYSTEM_STORE_PUSH_TASK_LOG_PREFIX, topicName); + break; + } + Utils.sleep(TimeUnit.SECONDS.toMillis(JOB_POLLING_RETRY_PERIOD_IN_SECONDS)); + count--; + if (count == 0) { + LOGGER.error( + "{} Push not finished: {} in {} seconds", + SYSTEM_STORE_PUSH_TASK_LOG_PREFIX, + topicName, + JOB_POLLING_RETRY_COUNT * JOB_POLLING_RETRY_PERIOD_IN_SECONDS); + return false; + } + } + } + return true; + } + + int getStoreLargestUsedVersionNumber(ControllerClient controllerClient, String systemStoreName) { + // Make sure store exist in region and return largest used version number. + StoreResponse systemStoreResponse = controllerClient.getStore(systemStoreName); + if (systemStoreResponse.isError()) { + return -1; + } + return systemStoreResponse.getStore().getLargestUsedVersionNumber(); + } +} diff --git a/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/transport/GrpcTransportClient.java b/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/transport/GrpcTransportClient.java index 30029701f24..145fd370dae 100644 --- a/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/transport/GrpcTransportClient.java +++ b/clients/venice-client/src/main/java/com/linkedin/venice/fastclient/transport/GrpcTransportClient.java @@ -21,10 +21,8 @@ import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap; import io.grpc.ChannelCredentials; import io.grpc.Grpc; -import io.grpc.InsecureChannelCredentials; import io.grpc.ManagedChannel; import io.grpc.Status; -import io.grpc.TlsChannelCredentials; import io.grpc.stub.StreamObserver; import java.io.IOException; import java.util.Arrays; @@ -74,7 +72,7 @@ public GrpcTransportClient(GrpcClientConfig grpcClientConfig) { this.port = port; this.serverGrpcChannels = new VeniceConcurrentHashMap<>(); this.stubCache = new VeniceConcurrentHashMap<>(); - this.channelCredentials = buildChannelCredentials(sslFactory); + this.channelCredentials = GrpcUtils.buildChannelCredentials(sslFactory); } @Override @@ -99,25 +97,6 @@ public void close() throws IOException { r2TransportClientForNonStorageOps.close(); } - @VisibleForTesting - ChannelCredentials buildChannelCredentials(SSLFactory sslFactory) { - // TODO: Evaluate if this needs to fail instead since it depends on plain text support on server - if (sslFactory == null) { - return InsecureChannelCredentials.create(); - } - - try { - TlsChannelCredentials.Builder tlsBuilder = TlsChannelCredentials.newBuilder() - .keyManager(GrpcUtils.getKeyManagers(sslFactory)) - .trustManager(GrpcUtils.getTrustManagers(sslFactory)); - return tlsBuilder.build(); - } catch (Exception e) { - throw new VeniceClientException( - "Failed to initialize SSL channel credentials for Venice gRPC Transport Client", - e); - } - } - @VisibleForTesting VeniceClientRequest buildVeniceClientRequest(String[] requestParts, byte[] requestBody, boolean isSingleGet) { VeniceClientRequest.Builder requestBuilder = VeniceClientRequest.newBuilder() diff --git a/clients/venice-client/src/test/java/com/linkedin/venice/fastclient/transport/GrpcTransportClientTest.java b/clients/venice-client/src/test/java/com/linkedin/venice/fastclient/transport/GrpcTransportClientTest.java index 02b228ce1b3..98577fad8e1 100644 --- a/clients/venice-client/src/test/java/com/linkedin/venice/fastclient/transport/GrpcTransportClientTest.java +++ b/clients/venice-client/src/test/java/com/linkedin/venice/fastclient/transport/GrpcTransportClientTest.java @@ -1,20 +1,27 @@ package com.linkedin.venice.fastclient.transport; -import static org.mockito.Mockito.*; -import static org.testng.Assert.*; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.anyBoolean; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; import com.google.common.collect.ImmutableMap; import com.linkedin.r2.transport.common.Client; import com.linkedin.venice.HttpMethod; -import com.linkedin.venice.client.exceptions.VeniceClientException; import com.linkedin.venice.client.store.transport.TransportClient; import com.linkedin.venice.client.store.transport.TransportClientResponse; import com.linkedin.venice.fastclient.GrpcClientConfig; import com.linkedin.venice.protocols.VeniceClientRequest; import com.linkedin.venice.protocols.VeniceReadServiceGrpc; import com.linkedin.venice.protocols.VeniceServerResponse; -import com.linkedin.venice.security.SSLFactory; -import io.grpc.ChannelCredentials; import java.util.Collections; import java.util.Map; import java.util.concurrent.CompletableFuture; @@ -61,14 +68,6 @@ public void setUp() { grpcTransportClient = new GrpcTransportClient(mockClientConfig); } - @Test(expectedExceptions = VeniceClientException.class) - public void testBuildChannelCredentials() { - ChannelCredentials actualChannelCredentials = grpcTransportClient.buildChannelCredentials(null); - assertNotNull(actualChannelCredentials, "Null ssl factory should default to insecure channel credentials"); - - grpcTransportClient.buildChannelCredentials(mock(SSLFactory.class)); - } - @Test public void testBuildVeniceClientRequestForSingleGet() { VeniceClientRequest clientRequest = diff --git a/clients/venice-push-job/src/test/java/com/linkedin/venice/heartbeat/TestPushJobHeartbeatSender.java b/clients/venice-push-job/src/test/java/com/linkedin/venice/heartbeat/TestPushJobHeartbeatSender.java index 1b098e0217e..6a25bb9a758 100644 --- a/clients/venice-push-job/src/test/java/com/linkedin/venice/heartbeat/TestPushJobHeartbeatSender.java +++ b/clients/venice-push-job/src/test/java/com/linkedin/venice/heartbeat/TestPushJobHeartbeatSender.java @@ -1,6 +1,9 @@ package com.linkedin.venice.heartbeat; -import static org.mockito.Mockito.*; +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import com.linkedin.venice.controllerapi.ControllerClient; import com.linkedin.venice.controllerapi.MultiSchemaResponse; diff --git a/docs/quickstart/quickstart-multi-datacenter.md b/docs/quickstart/quickstart-multi-datacenter.md index 4f477314a01..cff057de15e 100644 --- a/docs/quickstart/quickstart-multi-datacenter.md +++ b/docs/quickstart/quickstart-multi-datacenter.md @@ -13,7 +13,7 @@ Follow this guide to set up a multi-datacenter venice cluster using docker image provided by Venice team. -#### Step 1: Install and set up Docker Engine +#### Step 1: Install and set up Docker Engine and docker-compose Follow https://docs.docker.com/engine/install/ to install docker and start docker engine diff --git a/docs/quickstart/quickstart-single-datacenter.md b/docs/quickstart/quickstart-single-datacenter.md index 20c8c1c1f4f..dbf16a7419b 100644 --- a/docs/quickstart/quickstart-single-datacenter.md +++ b/docs/quickstart/quickstart-single-datacenter.md @@ -14,7 +14,7 @@ Follow this guide to set up a simple venice cluster using docker images provided by Venice team. -#### Step 1: Install and set up Docker Engine +#### Step 1: Install and set up Docker Engine and docker-compose Follow https://docs.docker.com/engine/install/ to install docker and start docker engine @@ -30,7 +30,7 @@ Once containers are up and running, it will create a test cluster, namely, `veni Note: Make sure the `docker-compose-single-dc-setup.yaml` downloaded in step 2 is in the same directory from which you will run the following command. ``` -docker compose -f docker-compose-single-dc-setup.yaml up -d +docker-compose -f docker-compose-single-dc-setup.yaml up -d ``` #### Step 4: Access `venice-client` container's bash shell @@ -60,7 +60,7 @@ value schema: ``` Let's create a venice store: -``` +```bash ./create-store.sh http://venice-controller:5555 venice-cluster0 test-store sample-data/schema/keySchema.avsc sample-data/schema/valueSchema.avsc ``` @@ -72,8 +72,15 @@ key: 1 to 100 value: val1 to val100 ``` -Let's push the data: +##### Print dataset +```bash +./avro-to-json.sh sample-data/batch-push-data/kv_records.avro ``` + +##### Run a push job + +Let's push the data: +```bash ./run-vpj.sh sample-data/single-dc-configs/batch-push-job.properties ``` @@ -102,15 +109,21 @@ value=null #### Step 8: Update and add some new records using Incremental Push Venice supports incremental push which allows us to update values of existing rows or to add new rows in an existing store. In this example, we will -1. update values for keys from `51-100`. For example, the new value of `100` will be `val100_v1` -2. add new rows (key: `101-150`) +1. update values for keys from `91-100`. For example, the new value of `100` will be `val100_v1` +2. add new rows (key: `101-110`) +##### Print records to be updated and added to the existing dataset in the store +```bash +./avro-to-json.sh sample-data/inc-push-data/kv_records_v1.avro +``` + +##### Run incremental push job ```bash ./run-vpj.sh sample-data/single-dc-configs/inc-push-job.properties ``` #### Step 9: Read data from the store after Incremental Push -Incremental Push updated the values of keys 51-100 and added new rows 101-150. +Incremental Push updated the values of keys 91-100 and added new rows 101-110. Let's read the data once again. ```bash @@ -132,14 +145,15 @@ value=val101 #### Step 10: Exit `venice-client` -``` +```bash +# type exit command on the terminal or use cntrl + c exit ``` #### Step 11: Stop docker Tear down the venice cluster -``` -docker compose -f docker-compose-single-dc-setup.yaml down +```bash +docker-compose -f docker-compose-single-dc-setup.yaml down ``` ## Next steps diff --git a/gradle/spotbugs/exclude.xml b/gradle/spotbugs/exclude.xml index 3d5a8dbf544..2543f656ba1 100644 --- a/gradle/spotbugs/exclude.xml +++ b/gradle/spotbugs/exclude.xml @@ -76,6 +76,8 @@ + + @@ -481,8 +483,10 @@ - - + + + + @@ -494,4 +498,22 @@ + + + + + + + + + + + + + + + + + + diff --git a/integrations/venice-duckdb/build.gradle b/integrations/venice-duckdb/build.gradle new file mode 100644 index 00000000000..d49bd238dc1 --- /dev/null +++ b/integrations/venice-duckdb/build.gradle @@ -0,0 +1,18 @@ +dependencies { + implementation libraries.avro + implementation libraries.avroUtilCompatHelper + implementation libraries.duckdbJdbc + api libraries.log4j2api + + implementation project(':clients:da-vinci-client') + implementation project(':internal:venice-client-common') + + implementation project(':internal:venice-common') +} + +checkerFramework { + extraJavacArgs = ['-Xmaxerrs', '256'] + checkers = ['org.checkerframework.checker.nullness.NullnessChecker'] + skipCheckerFramework = true + excludeTests = true +} \ No newline at end of file diff --git a/integrations/venice-duckdb/src/main/java/com/linkedin/venice/duckdb/DuckDBDaVinciRecordTransformer.java b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/duckdb/DuckDBDaVinciRecordTransformer.java new file mode 100644 index 00000000000..5215dfb80ec --- /dev/null +++ b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/duckdb/DuckDBDaVinciRecordTransformer.java @@ -0,0 +1,161 @@ +package com.linkedin.venice.duckdb; + +import static com.linkedin.venice.sql.AvroToSQL.UnsupportedTypeHandling.FAIL; + +import com.linkedin.davinci.client.DaVinciRecordTransformer; +import com.linkedin.davinci.client.DaVinciRecordTransformerResult; +import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.sql.AvroToSQL; +import com.linkedin.venice.sql.InsertProcessor; +import com.linkedin.venice.sql.SQLUtils; +import com.linkedin.venice.sql.TableDefinition; +import com.linkedin.venice.utils.lazy.Lazy; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Set; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericRecord; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +public class DuckDBDaVinciRecordTransformer + extends DaVinciRecordTransformer { + private static final Logger LOGGER = LogManager.getLogger(DuckDBDaVinciRecordTransformer.class); + private static final String duckDBFilePath = "my_database.duckdb"; + private static final String deleteStatementTemplate = "DELETE FROM %s WHERE %s = ?;"; + private static final String createViewStatementTemplate = + "CREATE OR REPLACE VIEW current_version AS SELECT * FROM %s;"; + private static final String dropTableStatementTemplate = "DROP TABLE %s;"; + private final String storeNameWithoutVersionInfo; + private final String versionTableName; + private final String duckDBUrl; + private final Set columnsToProject; + + public DuckDBDaVinciRecordTransformer( + int storeVersion, + Schema keySchema, + Schema inputValueSchema, + Schema outputValueSchema, + boolean storeRecordsInDaVinci, + String baseDir, + String storeNameWithoutVersionInfo, + Set columnsToProject) { + super(storeVersion, keySchema, inputValueSchema, outputValueSchema, storeRecordsInDaVinci); + this.storeNameWithoutVersionInfo = storeNameWithoutVersionInfo; + this.versionTableName = buildStoreNameWithVersion(storeVersion); + this.duckDBUrl = "jdbc:duckdb:" + baseDir + "/" + duckDBFilePath; + this.columnsToProject = columnsToProject; + } + + @Override + public DaVinciRecordTransformerResult transform(Lazy key, Lazy value) { + // Record transformation happens inside processPut as we need access to the connection object to create the prepared + // statement + return new DaVinciRecordTransformerResult<>(DaVinciRecordTransformerResult.Result.UNCHANGED); + } + + @Override + public void processPut(Lazy key, Lazy value) { + // TODO: Pre-allocate the upsert statement and everything that goes into it, as much as possible. + Schema keySchema = key.get().getSchema(); + Schema valueSchema = value.get().getSchema(); + String upsertStatement = AvroToSQL.upsertStatement(versionTableName, keySchema, valueSchema, this.columnsToProject); + + // ToDo: Instead of creating a connection on every call, have a long-term connection. Maybe a connection pool? + try (Connection connection = DriverManager.getConnection(duckDBUrl)) { + // TODO: Pre-allocate the upsert processor as well + InsertProcessor upsertProcessor = AvroToSQL.upsertProcessor(keySchema, valueSchema, this.columnsToProject); + + // TODO: Pre-allocate the prepared statement (consider thread-local if it's not thread safe) + try (PreparedStatement preparedStatement = connection.prepareStatement(upsertStatement)) { + upsertProcessor.process(key.get(), value.get(), preparedStatement); + } + } catch (SQLException e) { + throw new RuntimeException(e); + } + } + + @Override + public void processDelete(Lazy key) { + // Unable to convert to prepared statement as table and column names can't be parameterized + // ToDo make delete non-hardcoded on primaryKey + String deleteStatement = String.format(deleteStatementTemplate, versionTableName, "key"); + + // ToDo: Instead of creating a connection on every call, have a long-term connection. Maybe a connection pool? + try (Connection connection = DriverManager.getConnection(duckDBUrl); + PreparedStatement stmt = connection.prepareStatement(deleteStatement)) { + stmt.setString(1, key.get().get("key").toString()); + stmt.execute(); + } catch (SQLException e) { + throw new RuntimeException(e); + } + } + + @Override + public void onStartVersionIngestion(boolean isCurrentVersion) { + try (Connection connection = DriverManager.getConnection(duckDBUrl); + Statement stmt = connection.createStatement()) { + TableDefinition desiredTableDefinition = AvroToSQL.getTableDefinition( + this.versionTableName, + getKeySchema(), + getOutputValueSchema(), + this.columnsToProject, + FAIL, + true); + TableDefinition existingTableDefinition = SQLUtils.getTableDefinition(this.versionTableName, connection); + if (existingTableDefinition == null) { + LOGGER.info("Table '{}' not found on disk, will create it from scratch", this.versionTableName); + String createTableStatement = SQLUtils.createTableStatement(desiredTableDefinition); + stmt.execute(createTableStatement); + } else if (existingTableDefinition.equals(desiredTableDefinition)) { + LOGGER.info("Table '{}' found on disk and its schema is compatible. Will reuse.", this.versionTableName); + } else { + // TODO: Handle the wiping and re-bootstrap automatically. + throw new VeniceException( + "Table '" + this.versionTableName + "' found on disk, but its schema is incompatible. Please wipe."); + } + + if (isCurrentVersion) { + // Unable to convert to prepared statement as table and column names can't be parameterized + String createViewStatement = String.format(createViewStatementTemplate, versionTableName); + stmt.execute(createViewStatement); + } + } catch (SQLException e) { + throw new RuntimeException(e); + } + } + + @Override + public void onEndVersionIngestion(int currentVersion) { + try (Connection connection = DriverManager.getConnection(duckDBUrl); + Statement stmt = connection.createStatement()) { + // Swap to current version + String currentVersionTableName = buildStoreNameWithVersion(currentVersion); + String createViewStatement = String.format(createViewStatementTemplate, currentVersionTableName); + stmt.execute(createViewStatement); + + if (currentVersion != getStoreVersion()) { + // Only drop non-current versions, e.g., the backup version getting retired. + + // Unable to convert to prepared statement as table and column names can't be parameterized + // Drop DuckDB table for storeVersion as it's retired + String dropTableStatement = String.format(dropTableStatementTemplate, versionTableName); + stmt.execute(dropTableStatement); + } + } catch (SQLException e) { + throw new RuntimeException(e); + } + } + + public String getDuckDBUrl() { + return duckDBUrl; + } + + public String buildStoreNameWithVersion(int version) { + return storeNameWithoutVersionInfo + "_v" + version; + } +} diff --git a/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/AvroToSQL.java b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/AvroToSQL.java new file mode 100644 index 00000000000..f0cb0aeca8e --- /dev/null +++ b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/AvroToSQL.java @@ -0,0 +1,193 @@ +package com.linkedin.venice.sql; + +import java.sql.JDBCType; +import java.util.ArrayList; +import java.util.Collections; +import java.util.EnumMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Set; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import org.apache.avro.Schema; + + +/** + * Utility intended to convert Avro -> SQL, including DDL and DML statements. + * + * Initially, this implementation may have a DuckDB slant, though in the long-run it should ideally be vendor-neutral. + */ +public class AvroToSQL { + public enum UnsupportedTypeHandling { + FAIL, SKIP; + } + + private static final Map AVRO_TO_JDBC_TYPE_MAPPING; + + static { + Map avroToJdbc = new EnumMap<>(Schema.Type.class); + + // avroToJdbc.put(Schema.Type.UNION, JDBCType.?); // Unions need special handling, see below + avroToJdbc.put(Schema.Type.FIXED, JDBCType.BINARY); + avroToJdbc.put(Schema.Type.STRING, JDBCType.VARCHAR); + avroToJdbc.put(Schema.Type.BYTES, JDBCType.VARBINARY); + avroToJdbc.put(Schema.Type.INT, JDBCType.INTEGER); + avroToJdbc.put(Schema.Type.LONG, JDBCType.BIGINT); + avroToJdbc.put(Schema.Type.FLOAT, JDBCType.FLOAT); + avroToJdbc.put(Schema.Type.DOUBLE, JDBCType.DOUBLE); + avroToJdbc.put(Schema.Type.BOOLEAN, JDBCType.BOOLEAN); + avroToJdbc.put(Schema.Type.NULL, JDBCType.NULL); + + // Unsupported for now, but eventually might be: + // avroToJdbc.put(Schema.Type.RECORD, JDBCType.STRUCT); + // avroToJdbc.put(Schema.Type.ENUM, JDBCType.?); + // avroToJdbc.put(Schema.Type.ARRAY, JDBCType.ARRAY); + // avroToJdbc.put(Schema.Type.MAP, JDBCType.?); + + AVRO_TO_JDBC_TYPE_MAPPING = Collections.unmodifiableMap(avroToJdbc); + } + + private AvroToSQL() { + /** + * Static util. + * + * N.B.: For now, this is fine. But later on, we may want to specialize some of the behavior for different DB + * vendors (e.g., to support both DuckDB and SQLite, or even others). At that point, we would likely want to + * leverage subclasses, and therefore it may be cleaner to make this class abstract and instantiable. That is + * fine, we'll cross that bridge when we get to it. + */ + } + + @Nonnull + public static TableDefinition getTableDefinition( + @Nonnull String tableName, + @Nonnull Schema keySchema, + @Nonnull Schema valueSchema, + @Nonnull Set columnsToProject, + @Nonnull UnsupportedTypeHandling unsupportedTypeHandling, + boolean primaryKey) { + List columnDefinitions = new ArrayList<>(); + int jdbcIndex = 1; + for (Schema.Field field: combineColumns(keySchema, valueSchema, columnsToProject)) { + JDBCType correspondingType = getCorrespondingType(field); + if (correspondingType == null) { + switch (unsupportedTypeHandling) { + case SKIP: + continue; + case FAIL: + Schema fieldSchema = field.schema(); + Schema.Type fieldType = fieldSchema.getType(); + throw new IllegalArgumentException(fieldType + " is not supported!"); + default: + // Defensive code (unreachable) + throw new IllegalStateException("Missing enum branch handling!"); + } + } + + boolean isPrimaryKey = primaryKey && keySchema.getFields().contains(field); + columnDefinitions.add( + new ColumnDefinition( + SQLUtils.cleanColumnName(field.name()), + correspondingType, + true, // TODO: plug nullability + isPrimaryKey ? IndexType.PRIMARY_KEY : null, + null, // TODO: plug default (if necessary)... + null, + jdbcIndex++)); + } + + return new TableDefinition(tableName, columnDefinitions); + } + + @Nonnull + public static String upsertStatement( + @Nonnull String tableName, + @Nonnull Schema keySchema, + @Nonnull Schema valueSchema, + @Nonnull Set columnsToProject) { + Set allColumns = combineColumns(keySchema, valueSchema, columnsToProject); + StringBuffer stringBuffer = new StringBuffer(); + stringBuffer.append("INSERT OR REPLACE INTO " + SQLUtils.cleanTableName(tableName) + " VALUES ("); + boolean firstColumn = true; + + for (Schema.Field field: allColumns) { + JDBCType correspondingType = getCorrespondingType(field); + if (correspondingType == null) { + // Skipped field. + continue; + } + + if (firstColumn) { + firstColumn = false; + } else { + stringBuffer.append(", "); + } + + stringBuffer.append("?"); + } + stringBuffer.append(");"); + + return stringBuffer.toString(); + } + + @Nonnull + public static InsertProcessor upsertProcessor( + @Nonnull Schema keySchema, + @Nonnull Schema valueSchema, + @Nonnull Set columnsToProject) { + return new InsertProcessor(keySchema, valueSchema, columnsToProject); + } + + @Nullable + static JDBCType getCorrespondingType(Schema.Field field) { + Schema fieldSchema = field.schema(); + Schema.Type fieldType = fieldSchema.getType(); + + // Unpack unions + if (fieldType == Schema.Type.UNION) { + List unionBranches = fieldSchema.getTypes(); + if (unionBranches.size() == 2) { + if (unionBranches.get(0).getType() == Schema.Type.NULL) { + fieldType = unionBranches.get(1).getType(); + } else if (unionBranches.get(1).getType() == Schema.Type.NULL) { + fieldType = unionBranches.get(0).getType(); + } else { + return null; + } + } else { + return null; + } + } + + return AVRO_TO_JDBC_TYPE_MAPPING.get(fieldType); + } + + @Nonnull + static Set combineColumns( + @Nonnull Schema keySchema, + @Nonnull Schema valueSchema, + @Nonnull Set columnsToProject) { + Objects.requireNonNull(keySchema); + Objects.requireNonNull(valueSchema); + Objects.requireNonNull(columnsToProject); + if (keySchema.getType() != Schema.Type.RECORD || valueSchema.getType() != Schema.Type.RECORD) { + // TODO: We can improve this to handle primitive types which aren't wrapped inside records. + throw new IllegalArgumentException("Only Avro records can have a corresponding CREATE TABLE statement."); + } + Set allColumns = new LinkedHashSet<>(keySchema.getFields().size() + valueSchema.getFields().size()); + allColumns.addAll(keySchema.getFields()); + for (Schema.Field field: valueSchema.getFields()) { + if (columnsToProject.isEmpty() || columnsToProject.contains(field.name())) { + if (!allColumns.add(field)) { + throw new IllegalArgumentException( + "The value field '" + field.name() + "' is also present in the key schema! " + + "Field names must not conflict across both key and value. " + + "This can be side-stepped by populating the columnsToProject param to include only unique fields."); + } + } + } + return allColumns; + } +} diff --git a/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/ColumnDefinition.java b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/ColumnDefinition.java new file mode 100644 index 00000000000..4dd5dc54b4e --- /dev/null +++ b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/ColumnDefinition.java @@ -0,0 +1,109 @@ +package com.linkedin.venice.sql; + +import java.sql.JDBCType; +import java.util.Objects; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + + +public class ColumnDefinition { + @Nonnull + private final String name; + @Nonnull + private final JDBCType type; + private final boolean nullable; + @Nullable + private final IndexType indexType; + @Nullable + private final String defaultValue; + @Nullable + private final String extra; + private final int jdbcIndex; + + public ColumnDefinition(@Nonnull String name, @Nonnull JDBCType type, int jdbcIndex) { + this(name, type, true, null, jdbcIndex); + } + + public ColumnDefinition( + @Nonnull String name, + @Nonnull JDBCType type, + boolean nullable, + @Nullable IndexType indexType, + int jdbcIndex) { + this(name, type, nullable, indexType, null, null, jdbcIndex); + } + + public ColumnDefinition( + @Nonnull String name, + @Nonnull JDBCType type, + boolean nullable, + @Nullable IndexType indexType, + @Nullable String defaultValue, + @Nullable String extra, + int jdbcIndex) { + this.name = Objects.requireNonNull(name); + this.type = Objects.requireNonNull(type); + this.nullable = nullable; + this.indexType = indexType; + this.defaultValue = defaultValue; + this.extra = extra; + this.jdbcIndex = jdbcIndex; + if (this.jdbcIndex < 1) { + throw new IllegalArgumentException("The jdbcIndex must be at least 1"); + } + } + + @Nonnull + public String getName() { + return name; + } + + @Nonnull + public JDBCType getType() { + return type; + } + + public boolean isNullable() { + return nullable; + } + + @Nullable + public IndexType getIndexType() { + return indexType; + } + + @Nullable + public String getDefaultValue() { + return defaultValue; + } + + @Nullable + public String getExtra() { + return extra; + } + + public int getJdbcIndex() { + return jdbcIndex; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + ColumnDefinition that = (ColumnDefinition) o; + + return this.nullable == that.nullable && this.jdbcIndex == that.jdbcIndex && this.name.equals(that.name) + && this.type == that.type && this.indexType == that.indexType && Objects.equals(defaultValue, that.defaultValue) + && Objects.equals(extra, that.extra); + } + + @Override + public int hashCode() { + return this.name.hashCode(); + } +} diff --git a/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/IndexType.java b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/IndexType.java new file mode 100644 index 00000000000..33cf64bb41a --- /dev/null +++ b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/IndexType.java @@ -0,0 +1,5 @@ +package com.linkedin.venice.sql; + +public enum IndexType { + PRIMARY_KEY, UNIQUE; +} diff --git a/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/InsertProcessor.java b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/InsertProcessor.java new file mode 100644 index 00000000000..af6aa0f2dd4 --- /dev/null +++ b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/InsertProcessor.java @@ -0,0 +1,200 @@ +package com.linkedin.venice.sql; + +import com.linkedin.venice.utils.ByteUtils; +import java.nio.ByteBuffer; +import java.sql.JDBCType; +import java.sql.PreparedStatement; +import java.sql.SQLException; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import java.util.Set; +import javax.annotation.Nonnull; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericRecord; + + +public class InsertProcessor { + private final int[] keyFieldIndexToJdbcIndexMapping; + private final int[] valueFieldIndexToJdbcIndexMapping; + private final int[] keyFieldIndexToUnionBranchIndex; + private final int[] valueFieldIndexToUnionBranchIndex; + private final JDBCType[] keyFieldIndexToCorrespondingType; + private final JDBCType[] valueFieldIndexToCorrespondingType; + + InsertProcessor(@Nonnull Schema keySchema, @Nonnull Schema valueSchema, @Nonnull Set columnsToProject) { + Objects.requireNonNull(keySchema); + Objects.requireNonNull(valueSchema); + Objects.requireNonNull(columnsToProject); + + int keyFieldCount = keySchema.getFields().size(); + this.keyFieldIndexToJdbcIndexMapping = new int[keyFieldCount]; + this.keyFieldIndexToUnionBranchIndex = new int[keyFieldCount]; + this.keyFieldIndexToCorrespondingType = new JDBCType[keyFieldCount]; + + int valueFieldCount = valueSchema.getFields().size(); + this.valueFieldIndexToJdbcIndexMapping = new int[valueFieldCount]; + this.valueFieldIndexToUnionBranchIndex = new int[valueFieldCount]; + this.valueFieldIndexToCorrespondingType = new JDBCType[valueFieldCount]; + + // N.B.: JDBC indices start at 1, not at 0. + int index = 1; + index = populateArrays( + index, + keySchema, + this.keyFieldIndexToJdbcIndexMapping, + this.keyFieldIndexToUnionBranchIndex, + this.keyFieldIndexToCorrespondingType, + Collections.emptySet()); // N.B.: All key columns must be projected. + populateArrays( + index, // N.B.: The same index value needs to carry over from key to value columns. + valueSchema, + this.valueFieldIndexToJdbcIndexMapping, + this.valueFieldIndexToUnionBranchIndex, + this.valueFieldIndexToCorrespondingType, + columnsToProject); + } + + public void process(GenericRecord key, GenericRecord value, PreparedStatement preparedStatement) { + try { + processRecord( + key, + preparedStatement, + this.keyFieldIndexToJdbcIndexMapping, + this.keyFieldIndexToUnionBranchIndex, + this.keyFieldIndexToCorrespondingType); + + processRecord( + value, + preparedStatement, + this.valueFieldIndexToJdbcIndexMapping, + this.valueFieldIndexToUnionBranchIndex, + this.valueFieldIndexToCorrespondingType); + + preparedStatement.execute(); + } catch (SQLException e) { + throw new RuntimeException(e); + } + } + + private int populateArrays( + int index, + @Nonnull Schema schema, + @Nonnull int[] avroFieldIndexToJdbcIndexMapping, + @Nonnull int[] avroFieldIndexToUnionBranchIndex, + @Nonnull JDBCType[] avroFieldIndexToCorrespondingType, + @Nonnull Set columnsToProject) { + for (Schema.Field field: schema.getFields()) { + JDBCType correspondingType = AvroToSQL.getCorrespondingType(field); + if (correspondingType == null) { + // Skipped field. + continue; + } + if (!columnsToProject.isEmpty() && !columnsToProject.contains(field.name())) { + // Column is not projected. + continue; + } + avroFieldIndexToJdbcIndexMapping[field.pos()] = index++; + avroFieldIndexToCorrespondingType[field.pos()] = correspondingType; + if (field.schema().getType() == Schema.Type.UNION) { + Schema fieldSchema = field.schema(); + List unionBranches = fieldSchema.getTypes(); + if (unionBranches.get(0).getType() == Schema.Type.NULL) { + avroFieldIndexToUnionBranchIndex[field.pos()] = 1; + } else if (unionBranches.get(1).getType() == Schema.Type.NULL) { + avroFieldIndexToUnionBranchIndex[field.pos()] = 0; + } else { + throw new IllegalStateException("Should have skipped unsupported union: " + fieldSchema); + } + } + } + return index; + } + + private void processRecord( + GenericRecord record, + PreparedStatement preparedStatement, + int[] avroFieldIndexToJdbcIndexMapping, + int[] avroFieldIndexToUnionBranchIndex, + JDBCType[] avroFieldIndexToCorrespondingType) throws SQLException { + int jdbcIndex; + JDBCType jdbcType; + Object fieldValue; + Schema.Type fieldType; + for (Schema.Field field: record.getSchema().getFields()) { + jdbcIndex = avroFieldIndexToJdbcIndexMapping[field.pos()]; + if (jdbcIndex == 0) { + // Skipped field. + continue; + } + fieldValue = record.get(field.pos()); + if (fieldValue == null) { + jdbcType = avroFieldIndexToCorrespondingType[field.pos()]; + preparedStatement.setNull(jdbcIndex, jdbcType.getVendorTypeNumber()); + continue; + } + fieldType = field.schema().getType(); + if (fieldType == Schema.Type.UNION) { + // Unions are handled via unpacking + fieldType = field.schema().getTypes().get(avroFieldIndexToUnionBranchIndex[field.pos()]).getType(); + } + try { + processField(jdbcIndex, fieldType, fieldValue, preparedStatement, field.name()); + } catch (Exception e) { + throw new RuntimeException( + "Failed to process field. Name: '" + field.name() + "; jdbcIndex: " + jdbcIndex + "; type: " + fieldType + + "; value: " + fieldValue, + e); + } + } + } + + private void processField( + int jdbcIndex, + @Nonnull Schema.Type fieldType, + @Nonnull Object fieldValue, + @Nonnull PreparedStatement preparedStatement, + @Nonnull String fieldName) throws SQLException { + switch (fieldType) { + case FIXED: + case BYTES: + preparedStatement.setBytes(jdbcIndex, ByteUtils.extractByteArray((ByteBuffer) fieldValue)); + break; + case STRING: + preparedStatement.setString(jdbcIndex, fieldValue.toString()); + break; + case INT: + preparedStatement.setInt(jdbcIndex, (int) fieldValue); + break; + case LONG: + preparedStatement.setLong(jdbcIndex, (long) fieldValue); + break; + case FLOAT: + preparedStatement.setFloat(jdbcIndex, (float) fieldValue); + break; + case DOUBLE: + preparedStatement.setDouble(jdbcIndex, (double) fieldValue); + break; + case BOOLEAN: + preparedStatement.setBoolean(jdbcIndex, (boolean) fieldValue); + break; + case NULL: + // Weird case... probably never comes into play? + preparedStatement.setNull(jdbcIndex, JDBCType.NULL.getVendorTypeNumber()); + break; + + case UNION: + // Defensive code. Unreachable. + throw new IllegalArgumentException( + "Unions should be unpacked by the calling function, but union field '" + fieldName + "' was passed in!"); + + // These types could be supported eventually, but for now aren't. + case RECORD: + case ENUM: + case ARRAY: + case MAP: + default: + throw new IllegalStateException("Should have skipped field '" + fieldName + "' but somehow didn't!"); + } + } +} diff --git a/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/SQLUtils.java b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/SQLUtils.java new file mode 100644 index 00000000000..f27c51f00cc --- /dev/null +++ b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/SQLUtils.java @@ -0,0 +1,130 @@ +package com.linkedin.venice.sql; + +import java.sql.Connection; +import java.sql.JDBCType; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.List; +import java.util.Objects; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + + +public class SQLUtils { + private static final String SHOW_TABLE_COLUMN_NAME = "column_name"; + private static final String SHOW_TABLE_TYPE = "column_type"; + private static final String SHOW_TABLE_NULLABLE = "null"; + private static final String SHOW_TABLE_INDEX_TYPE = "key"; + private static final String SHOW_TABLE_DEFAULT = "default"; + private static final String SHOW_TABLE_EXTRA = "extra"; + + private SQLUtils() { + /** + * Static utils. + * + * If we wish to specialize the behavior of this class for different DB vendors, then we can change it to + * an abstract class and have instantiable subclasses for each vendor. + */ + } + + @Nullable + public static TableDefinition getTableDefinition(@Nonnull String tableName, Connection connection) + throws SQLException { + try (PreparedStatement preparedStatement = + connection.prepareStatement("SELECT * FROM (SHOW TABLES) WHERE name = ?;")) { + preparedStatement.setString(1, cleanTableName(tableName)); + try (ResultSet resultSet = preparedStatement.executeQuery()) { + if (!resultSet.next()) { + return null; + } + } + } + try (Statement statement = connection.createStatement(); + ResultSet resultSet = statement.executeQuery("SHOW " + cleanTableName(tableName))) { + int jdbcIndex = 1; + List columns = new ArrayList<>(); + while (resultSet.next()) { + columns.add(convertToColumnDefinition(resultSet, jdbcIndex++)); + } + return new TableDefinition(tableName, columns); + } + } + + @Nonnull + public static String createTableStatement(TableDefinition tableDefinition) { + StringBuffer stringBuffer = new StringBuffer(); + stringBuffer.append("CREATE TABLE " + cleanTableName(tableDefinition.getName()) + "("); + boolean firstColumn = true; + + for (ColumnDefinition columnDefinition: tableDefinition.getColumns()) { + if (firstColumn) { + firstColumn = false; + } else { + stringBuffer.append(", "); + } + + stringBuffer.append(cleanColumnName(columnDefinition.getName()) + " " + columnDefinition.getType().name()); + } + + firstColumn = true; + if (!tableDefinition.getPrimaryKeyColumns().isEmpty()) { + stringBuffer.append(", PRIMARY KEY("); + for (ColumnDefinition columnDefinition: tableDefinition.getPrimaryKeyColumns()) { + if (firstColumn) { + firstColumn = false; + } else { + stringBuffer.append(", "); + } + stringBuffer.append(cleanColumnName(columnDefinition.getName())); + } + stringBuffer.append(")"); + } + stringBuffer.append(");"); + + return stringBuffer.toString(); + } + + /** + * This function should encapsulate the handling of any illegal characters (by either failing or converting them). + */ + @Nonnull + static String cleanTableName(@Nonnull String avroRecordName) { + return Objects.requireNonNull(avroRecordName); + } + + /** + * This function should encapsulate the handling of any illegal characters (by either failing or converting them). + */ + @Nonnull + static String cleanColumnName(@Nonnull String avroFieldName) { + return Objects.requireNonNull(avroFieldName); + } + + /** + * {@link ResultSet#next()} should have already been called upstream. This function will not call it. + */ + private static ColumnDefinition convertToColumnDefinition(ResultSet resultSet, int jdbcIndex) throws SQLException { + return new ColumnDefinition( + resultSet.getString(SHOW_TABLE_COLUMN_NAME), + JDBCType.valueOf(resultSet.getString(SHOW_TABLE_TYPE)), + resultSet.getString(SHOW_TABLE_NULLABLE).equals("YES"), + convertIndexType(resultSet.getString(SHOW_TABLE_INDEX_TYPE)), + resultSet.getString(SHOW_TABLE_DEFAULT), + resultSet.getString(SHOW_TABLE_EXTRA), + jdbcIndex); + } + + private static IndexType convertIndexType(String indexType) { + if (indexType == null) { + return null; + } else if (indexType.equals("PRI")) { + return IndexType.PRIMARY_KEY; + } else if (indexType.equals("UNI")) { + return IndexType.UNIQUE; + } + throw new IllegalArgumentException("Unsupported index type: " + indexType); + } +} diff --git a/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/TableDefinition.java b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/TableDefinition.java new file mode 100644 index 00000000000..ef4710e233e --- /dev/null +++ b/integrations/venice-duckdb/src/main/java/com/linkedin/venice/sql/TableDefinition.java @@ -0,0 +1,102 @@ +package com.linkedin.venice.sql; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; + + +public class TableDefinition { + @Nonnull + private final String name; + @Nonnull + private final List columns; + @Nonnull + private final List primaryKeyColumns; + + public TableDefinition(@Nonnull String name, @Nonnull List columns) { + this.name = Objects.requireNonNull(name); + this.columns = Collections.unmodifiableList(columns); + if (this.columns.isEmpty()) { + throw new IllegalArgumentException("columns cannot be empty!"); + } + for (int i = 0; i < this.columns.size(); i++) { + ColumnDefinition columnDefinition = this.columns.get(i); + if (columnDefinition == null) { + throw new IllegalArgumentException( + "The columns list must be densely populated, but found a null entry at index: " + i); + } + int expectedJdbcIndex = i + 1; + if (columnDefinition.getJdbcIndex() != expectedJdbcIndex) { + throw new IllegalArgumentException( + "The columns list must be populated in the correct order, column '" + columnDefinition.getName() + + "' found at index " + i + " but its JDBC index is " + columnDefinition.getJdbcIndex() + " (expected " + + expectedJdbcIndex + ")."); + } + } + List pkColumns = new ArrayList<>(); + for (ColumnDefinition columnDefinition: this.columns) { + if (columnDefinition.getIndexType() == IndexType.PRIMARY_KEY) { + pkColumns.add(columnDefinition); + } + } + this.primaryKeyColumns = pkColumns.isEmpty() ? Collections.emptyList() : Collections.unmodifiableList(pkColumns); + } + + @Nonnull + public String getName() { + return name; + } + + @Nonnull + public List getColumns() { + return columns; + } + + @Nonnull + ColumnDefinition getColumnByJdbcIndex(int index) { + if (index < 1 || index > this.columns.size()) { + throw new IndexOutOfBoundsException("Invalid index. The valid range is: 1.." + this.columns.size()); + } + return this.columns.get(index - 1); + } + + @Nullable + ColumnDefinition getColumnByName(@Nonnull String columnName) { + Objects.requireNonNull(columnName); + for (int i = 0; i < this.columns.size(); i++) { + ColumnDefinition columnDefinition = this.columns.get(i); + if (columnDefinition.getName().equals(columnName)) { + return columnDefinition; + } + } + return null; + } + + @Nonnull + public List getPrimaryKeyColumns() { + return primaryKeyColumns; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (o == null || getClass() != o.getClass()) { + return false; + } + + TableDefinition that = (TableDefinition) o; + + return this.name.equals(that.name) && this.columns.equals(that.columns) + && this.primaryKeyColumns.equals(that.primaryKeyColumns); + } + + @Override + public int hashCode() { + return this.name.hashCode(); + } +} diff --git a/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBAvroToSQLTest.java b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBAvroToSQLTest.java new file mode 100644 index 00000000000..b5af04d373d --- /dev/null +++ b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBAvroToSQLTest.java @@ -0,0 +1,214 @@ +package com.linkedin.venice.duckdb; + +import static com.linkedin.avroutil1.compatibility.AvroCompatibilityHelper.createSchemaField; +import static com.linkedin.venice.sql.AvroToSQL.UnsupportedTypeHandling.SKIP; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + +import com.linkedin.venice.sql.AvroToSQL; +import com.linkedin.venice.sql.AvroToSQLTest; +import com.linkedin.venice.sql.InsertProcessor; +import com.linkedin.venice.sql.SQLUtils; +import com.linkedin.venice.utils.ByteUtils; +import java.io.IOException; +import java.nio.ByteBuffer; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.AbstractMap; +import java.util.ArrayList; +import java.util.Collections; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Random; +import java.util.Set; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.apache.commons.io.IOUtils; +import org.duckdb.DuckDBResultSet; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + + +public class DuckDBAvroToSQLTest { + @DataProvider + public Object[][] primaryKeyColumns() { + Set compositePK = new HashSet<>(); + compositePK.add("intField"); + compositePK.add("longField"); + return new Object[][] { { Collections.singleton("intField") }, { compositePK } }; + } + + @DataProvider + public Object[][] primaryKeySchemas() { + Schema singleColumnKey = Schema.createRecord( + "MyKey", + "", + "", + false, + Collections.singletonList(createSchemaField("key1", Schema.create(Schema.Type.INT), "", null))); + List compositeKeyFields = new ArrayList<>(); + compositeKeyFields.add(createSchemaField("key1", Schema.create(Schema.Type.INT), "", null)); + compositeKeyFields.add(createSchemaField("key2", Schema.create(Schema.Type.LONG), "", null)); + Schema compositeKey = Schema.createRecord("MyKey", "", "", false, compositeKeyFields); + return new Object[][] { { singleColumnKey }, { compositeKey } }; + } + + @Test(dataProvider = "primaryKeySchemas") + public void testUpsert(Schema keySchema) throws SQLException, IOException { + List fields = AvroToSQLTest.getAllValidFields(); + Schema valueSchema = Schema.createRecord("MyRecord", "", "", false, fields); + try (Connection connection = DriverManager.getConnection("jdbc:duckdb:"); + Statement stmt = connection.createStatement()) { + // create a table + String tableName = "MyRecord_v1"; + String createTableStatement = SQLUtils.createTableStatement( + AvroToSQL.getTableDefinition(tableName, keySchema, valueSchema, Collections.emptySet(), SKIP, true)); + System.out.println(createTableStatement); + stmt.execute(createTableStatement); + + String upsertStatement = AvroToSQL.upsertStatement(tableName, keySchema, valueSchema, Collections.emptySet()); + System.out.println(upsertStatement); + InsertProcessor upsertProcessor = AvroToSQL.upsertProcessor(keySchema, valueSchema, Collections.emptySet()); + for (int rewriteIteration = 0; rewriteIteration < 3; rewriteIteration++) { + List> records = generateRecords(keySchema, valueSchema); + GenericRecord key, value; + try (PreparedStatement preparedStatement = connection.prepareStatement(upsertStatement)) { + for (int i = 0; i < records.size(); i++) { + key = records.get(i).getKey(); + value = records.get(i).getValue(); + upsertProcessor.process(key, value, preparedStatement); + } + } + + int recordCount = records.size(); + try (ResultSet rs = stmt.executeQuery("SELECT * FROM " + tableName)) { + for (int j = 0; j < recordCount; j++) { + assertTrue( + rs.next(), + "Rewrite iteration " + rewriteIteration + ". Expected a row at index " + j + + " after having inserted up to row " + recordCount); + + key = records.get(j).getKey(); + value = records.get(j).getValue(); + assertRecordValidity(keySchema, key, rs, j); + assertRecordValidity(valueSchema, value, rs, j); + System.out.println("Rewrite iteration " + rewriteIteration + ". Successfully validated row " + j); + } + assertFalse(rs.next(), "Expected no more rows at index " + recordCount); + } + System.out + .println("Rewrite iteration " + rewriteIteration + ". Successfully validated up to i = " + recordCount); + } + } + } + + private void assertRecordValidity(Schema schema, GenericRecord record, ResultSet rs, int j) + throws IOException, SQLException { + for (Schema.Field field: schema.getFields()) { + Object result = rs.getObject(field.name()); + if (result instanceof DuckDBResultSet.DuckDBBlobResult) { + DuckDBResultSet.DuckDBBlobResult duckDBBlobResult = (DuckDBResultSet.DuckDBBlobResult) result; + byte[] actual = IOUtils.toByteArray(duckDBBlobResult.getBinaryStream()); + byte[] expected = ((ByteBuffer) record.get(field.name())).array(); + assertEquals( + actual, + expected, + ". Bytes not equals at row " + j + "! actual: " + ByteUtils.toHexString(actual) + ", wanted: " + + ByteUtils.toHexString(expected)); + // System.out.println("Rewrite iteration " + rewriteIteration + ". Row: " + j + ", field: " + + // field.name() + ", value: " + ByteUtils.toHexString(actual)); + } else { + assertEquals( + result, + record.get(field.name()), + ". Field '" + field.name() + "' is not correct at row " + j + "!"); + // System.out.println("Rewrite iteration " + rewriteIteration + ". Row: " + j + ", field: " + + // field.name() + ", value: " + result); + } + } + + } + + private List> generateRecords(Schema keySchema, Schema valueSchema) { + List> records = new ArrayList<>(); + + GenericRecord keyRecord, valueRecord; + Object fieldValue; + Random random = new Random(); + for (int i = 0; i < 10; i++) { + keyRecord = new GenericData.Record(keySchema); + for (Schema.Field field: keySchema.getFields()) { + switch (field.schema().getType()) { + case INT: + fieldValue = i; + break; + case LONG: + fieldValue = (long) i; + break; + default: + throw new IllegalArgumentException("Only numeric PK columns are supported in this test."); + } + keyRecord.put(field.pos(), fieldValue); + } + + valueRecord = new GenericData.Record(valueSchema); + for (Schema.Field field: valueSchema.getFields()) { + Schema fieldSchema = field.schema(); + if (fieldSchema.getType() == Schema.Type.UNION) { + Schema first = field.schema().getTypes().get(0); + Schema second = field.schema().getTypes().get(1); + if (first.getType() == Schema.Type.NULL) { + fieldSchema = second; + } else if (second.getType() == Schema.Type.NULL) { + fieldSchema = first; + } else { + throw new IllegalArgumentException("Unsupported union: " + field.schema()); + } + } + fieldValue = randomValue(fieldSchema, random); + valueRecord.put(field.pos(), fieldValue); + } + records.add(new AbstractMap.SimpleEntry<>(keyRecord, valueRecord)); + } + + return records; + } + + private Object randomValue(Schema schema, Random random) { + switch (schema.getType()) { + case STRING: + return String.valueOf(random.nextLong()); + case INT: + return random.nextInt(); + case LONG: + return random.nextLong(); + case FLOAT: + return random.nextFloat(); + case DOUBLE: + return random.nextDouble(); + case BOOLEAN: + return random.nextBoolean(); + case BYTES: + return getBB(10, random); + case FIXED: + return getBB(schema.getFixedSize(), random); + case NULL: + return null; + default: + throw new IllegalArgumentException("Unsupported type: " + schema.getType()); + } + } + + private ByteBuffer getBB(int size, Random random) { + byte[] bytes = new byte[size]; + random.nextBytes(bytes); + return ByteBuffer.wrap(bytes); + } +} diff --git a/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBDaVinciRecordTransformerTest.java b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBDaVinciRecordTransformerTest.java new file mode 100644 index 00000000000..3f6e7644dac --- /dev/null +++ b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBDaVinciRecordTransformerTest.java @@ -0,0 +1,158 @@ +package com.linkedin.venice.duckdb; + +import static com.linkedin.venice.utils.TestWriteUtils.NAME_RECORD_V1_SCHEMA; +import static com.linkedin.venice.utils.TestWriteUtils.SINGLE_FIELD_RECORD_SCHEMA; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; + +import com.linkedin.davinci.client.DaVinciRecordTransformerResult; +import com.linkedin.davinci.client.DaVinciRecordTransformerUtility; +import com.linkedin.venice.utils.Utils; +import com.linkedin.venice.utils.lazy.Lazy; +import java.io.File; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Collections; +import java.util.Set; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class DuckDBDaVinciRecordTransformerTest { + static final int storeVersion = 1; + static final String storeName = "test_store"; + private final Set columnsToProject = Collections.emptySet(); + + @BeforeMethod + @AfterClass + public void deleteClassHash() { + File file = new File(String.format("./classHash-%d.txt", storeVersion)); + if (file.exists()) { + assertTrue(file.delete()); + } + } + + @Test + public void testRecordTransformer() { + String tempDir = Utils.getTempDataDirectory().getAbsolutePath(); + + DuckDBDaVinciRecordTransformer recordTransformer = new DuckDBDaVinciRecordTransformer( + storeVersion, + SINGLE_FIELD_RECORD_SCHEMA, + NAME_RECORD_V1_SCHEMA, + NAME_RECORD_V1_SCHEMA, + false, + tempDir, + storeName, + columnsToProject); + + Schema keySchema = recordTransformer.getKeySchema(); + assertEquals(keySchema.getType(), Schema.Type.RECORD); + + Schema outputValueSchema = recordTransformer.getOutputValueSchema(); + assertEquals(outputValueSchema.getType(), Schema.Type.RECORD); + + recordTransformer.onStartVersionIngestion(true); + + GenericRecord keyRecord = new GenericData.Record(SINGLE_FIELD_RECORD_SCHEMA); + keyRecord.put("key", "key"); + Lazy lazyKey = Lazy.of(() -> keyRecord); + + GenericRecord valueRecord = new GenericData.Record(NAME_RECORD_V1_SCHEMA); + valueRecord.put("firstName", "Duck"); + valueRecord.put("lastName", "Goose"); + Lazy lazyValue = Lazy.of(() -> valueRecord); + + DaVinciRecordTransformerResult transformerResult = recordTransformer.transform(lazyKey, lazyValue); + recordTransformer.processPut(lazyKey, lazyValue); + assertEquals(transformerResult.getResult(), DaVinciRecordTransformerResult.Result.UNCHANGED); + // Result will be empty when it's UNCHANGED + assertNull(transformerResult.getValue()); + assertNull(recordTransformer.transformAndProcessPut(lazyKey, lazyValue)); + + recordTransformer.processDelete(lazyKey); + + assertFalse(recordTransformer.getStoreRecordsInDaVinci()); + + int classHash = recordTransformer.getClassHash(); + + DaVinciRecordTransformerUtility recordTransformerUtility = + recordTransformer.getRecordTransformerUtility(); + assertTrue(recordTransformerUtility.hasTransformerLogicChanged(classHash)); + assertFalse(recordTransformerUtility.hasTransformerLogicChanged(classHash)); + } + + @Test + public void testVersionSwap() throws SQLException { + String tempDir = Utils.getTempDataDirectory().getAbsolutePath(); + + DuckDBDaVinciRecordTransformer recordTransformer_v1 = new DuckDBDaVinciRecordTransformer( + 1, + SINGLE_FIELD_RECORD_SCHEMA, + NAME_RECORD_V1_SCHEMA, + NAME_RECORD_V1_SCHEMA, + false, + tempDir, + storeName, + columnsToProject); + DuckDBDaVinciRecordTransformer recordTransformer_v2 = new DuckDBDaVinciRecordTransformer( + 2, + SINGLE_FIELD_RECORD_SCHEMA, + NAME_RECORD_V1_SCHEMA, + NAME_RECORD_V1_SCHEMA, + false, + tempDir, + storeName, + columnsToProject); + + String duckDBUrl = recordTransformer_v1.getDuckDBUrl(); + + recordTransformer_v1.onStartVersionIngestion(true); + recordTransformer_v2.onStartVersionIngestion(false); + + GenericRecord keyRecord = new GenericData.Record(SINGLE_FIELD_RECORD_SCHEMA); + keyRecord.put("key", "key"); + Lazy lazyKey = Lazy.of(() -> keyRecord); + + GenericRecord valueRecord_v1 = new GenericData.Record(NAME_RECORD_V1_SCHEMA); + valueRecord_v1.put("firstName", "Duck"); + valueRecord_v1.put("lastName", "Goose"); + Lazy lazyValue = Lazy.of(() -> valueRecord_v1); + recordTransformer_v1.processPut(lazyKey, lazyValue); + + GenericRecord valueRecord_v2 = new GenericData.Record(NAME_RECORD_V1_SCHEMA); + valueRecord_v2.put("firstName", "Goose"); + valueRecord_v2.put("lastName", "Duck"); + lazyValue = Lazy.of(() -> valueRecord_v2); + recordTransformer_v2.processPut(lazyKey, lazyValue); + + try (Connection connection = DriverManager.getConnection(duckDBUrl); + Statement stmt = connection.createStatement()) { + try (ResultSet rs = stmt.executeQuery("SELECT * FROM current_version")) { + assertTrue(rs.next(), "There should be a first row!"); + assertEquals(rs.getString("firstName"), "Duck"); + assertEquals(rs.getString("lastName"), "Goose"); + assertFalse(rs.next(), "There should be only one row!"); + } + + // Swap here + recordTransformer_v1.onEndVersionIngestion(2); + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM current_version")) { + assertTrue(rs.next(), "There should be a first row!"); + assertEquals(rs.getString("firstName"), "Goose"); + assertEquals(rs.getString("lastName"), "Duck"); + assertFalse(rs.next(), "There should be only one row!"); + } + } + } +} diff --git a/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBHelloWorldTest.java b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBHelloWorldTest.java new file mode 100644 index 00000000000..64a9ec5fb2e --- /dev/null +++ b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBHelloWorldTest.java @@ -0,0 +1,59 @@ +package com.linkedin.venice.duckdb; + +import com.linkedin.venice.sql.SQLHelloWorldTest; +import com.linkedin.venice.utils.Utils; +import java.io.File; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import org.duckdb.DuckDBAppender; +import org.duckdb.DuckDBConnection; +import org.testng.annotations.Test; + + +/** + * The aim of this class is just to test DuckDB itself, without any Venice-ism involved. + */ +public class DuckDBHelloWorldTest extends SQLHelloWorldTest { + public static final String DUCKDB_CONN_PREFIX = "jdbc:duckdb:"; + + @Override + protected Connection getConnection() throws SQLException { + return DriverManager.getConnection(DUCKDB_CONN_PREFIX); + } + + @Override + protected String getConnectionStringToPersistentDB() { + File tmpDir = Utils.getTempDataDirectory(); + return DUCKDB_CONN_PREFIX + tmpDir.getAbsolutePath() + "/foo.duckdb"; + } + + @Test + public void testAppender() throws SQLException { + try (DuckDBConnection connection = (DuckDBConnection) getConnection(); + Statement stmt = connection.createStatement()) { + // create a table + stmt.execute(createTableStatement("items", true)); + // insert two items into the table + try (DuckDBAppender appender = connection.createAppender(DuckDBConnection.DEFAULT_SCHEMA, "items")) { + appender.beginRow(); + appender.append("jeans"); + appender.append(20.0); + appender.append(1); + appender.endRow(); + + appender.beginRow(); + appender.append("hammer"); + appender.append(42.2); + appender.append(2); + appender.endRow(); + } + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM items")) { + assertValidityOfResultSet1(rs); + } + } + } +} diff --git a/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBSQLUtilsTest.java b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBSQLUtilsTest.java new file mode 100644 index 00000000000..a2cf4319c4c --- /dev/null +++ b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/duckdb/DuckDBSQLUtilsTest.java @@ -0,0 +1,16 @@ +package com.linkedin.venice.duckdb; + +import com.linkedin.venice.sql.SQLUtilsTest; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.SQLException; +import org.testng.annotations.Test; + + +@Test +public class DuckDBSQLUtilsTest extends SQLUtilsTest { + @Override + protected Connection getConnection() throws SQLException { + return DriverManager.getConnection(DuckDBHelloWorldTest.DUCKDB_CONN_PREFIX); + } +} diff --git a/integrations/venice-duckdb/src/test/java/com/linkedin/venice/sql/AvroToSQLTest.java b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/sql/AvroToSQLTest.java new file mode 100644 index 00000000000..22a249978a7 --- /dev/null +++ b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/sql/AvroToSQLTest.java @@ -0,0 +1,235 @@ +package com.linkedin.venice.sql; + +import static com.linkedin.avroutil1.compatibility.AvroCompatibilityHelper.createSchemaField; +import static com.linkedin.venice.sql.AvroToSQL.UnsupportedTypeHandling.FAIL; +import static com.linkedin.venice.sql.AvroToSQL.UnsupportedTypeHandling.SKIP; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertThrows; + +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.apache.avro.Schema; +import org.testng.annotations.Test; + + +public class AvroToSQLTest { + private static final String EXPECTED_CREATE_TABLE_STATEMENT_WITH_ALL_TYPES = "CREATE TABLE MyRecord(" // + + "key1 INTEGER, " // + + "fixedField BINARY, " // + + "stringField VARCHAR, " // + + "bytesField VARBINARY, "// + + "intField INTEGER, " // + + "longField BIGINT, " // + + "floatField FLOAT, " // + + "doubleField DOUBLE, " // + + "booleanField BOOLEAN, " // + // + "nullField NULL, " // + + "fixedFieldUnion1 BINARY, " // + + "fixedFieldUnion2 BINARY, " // + + "stringFieldUnion1 VARCHAR, " // + + "stringFieldUnion2 VARCHAR, " // + + "bytesFieldUnion1 VARBINARY, " // + + "bytesFieldUnion2 VARBINARY, " // + + "intFieldUnion1 INTEGER, " // + + "intFieldUnion2 INTEGER, " // + + "longFieldUnion1 BIGINT, " // + + "longFieldUnion2 BIGINT, " // + + "floatFieldUnion1 FLOAT, " // + + "floatFieldUnion2 FLOAT, " // + + "doubleFieldUnion1 DOUBLE, " // + + "doubleFieldUnion2 DOUBLE, " // + + "booleanFieldUnion1 BOOLEAN, " // + + "booleanFieldUnion2 BOOLEAN);"; + + private static final String EXPECTED_UPSERT_STATEMENT_WITH_ALL_TYPES = + "INSERT OR REPLACE INTO MyRecord VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?);"; + + @Test + public void testValidCreateTable() { + List allFields = getAllValidFields(); + Schema schemaWithAllSupportedFieldTypes = Schema.createRecord("MyRecord", "", "", false, allFields); + + // Single-column key (but without a PRIMARY KEY constraint) + Schema singleColumnKey = Schema.createRecord("MyKey", "", "", false, Collections.singletonList(getKey1())); + + TableDefinition tableDefinitionWithoutPrimaryKey = AvroToSQL.getTableDefinition( + "MyRecord", + singleColumnKey, + schemaWithAllSupportedFieldTypes, + Collections.emptySet(), + FAIL, + false); + String createTableWithoutPrimaryKey = SQLUtils.createTableStatement(tableDefinitionWithoutPrimaryKey); + assertEquals(createTableWithoutPrimaryKey, EXPECTED_CREATE_TABLE_STATEMENT_WITH_ALL_TYPES); + + // Single-column primary key + TableDefinition tableDefinitionWithPrimaryKey = AvroToSQL.getTableDefinition( + "MyRecord", + singleColumnKey, + schemaWithAllSupportedFieldTypes, + Collections.emptySet(), + FAIL, + true); + String createTableWithPrimaryKey = SQLUtils.createTableStatement(tableDefinitionWithPrimaryKey); + String expectedCreateTable = EXPECTED_CREATE_TABLE_STATEMENT_WITH_ALL_TYPES.replace(");", ", PRIMARY KEY(key1));"); + assertEquals(createTableWithPrimaryKey, expectedCreateTable); + + // Composite primary key + List compositeKeyFields = new ArrayList<>(); + compositeKeyFields.add(getKey1()); + compositeKeyFields.add(createSchemaField("key2", Schema.create(Schema.Type.LONG), "", null)); + Schema compositeKey = Schema.createRecord("MyKey", "", "", false, compositeKeyFields); + + TableDefinition tableDefinitionWithCompositePrimaryKey = AvroToSQL.getTableDefinition( + "MyRecord", + compositeKey, + schemaWithAllSupportedFieldTypes, + Collections.emptySet(), + FAIL, + true); + String createTableWithCompositePrimaryKey = SQLUtils.createTableStatement(tableDefinitionWithCompositePrimaryKey); + String expectedCreateTableWithCompositePK = + EXPECTED_CREATE_TABLE_STATEMENT_WITH_ALL_TYPES.replace("key1 INTEGER", "key1 INTEGER, key2 BIGINT") + .replace(");", ", PRIMARY KEY(key1, key2));"); + assertEquals(createTableWithCompositePrimaryKey, expectedCreateTableWithCompositePK); + } + + @Test + public void testUnsupportedTypesHandling() { + // Types that will for sure not be supported. + + assertThrows( + IllegalArgumentException.class, + () -> SQLUtils.createTableStatement( + AvroToSQL.getTableDefinition( + "MyRecord", + Schema.create(Schema.Type.INT), + Schema.create(Schema.Type.INT), + Collections.emptySet(), + FAIL, + true))); + + testSchemaWithInvalidType( + createSchemaField( + "TripleUnionWithNull", + Schema.createUnion( + Schema.create(Schema.Type.NULL), + Schema.create(Schema.Type.INT), + Schema.create(Schema.Type.STRING)), + "", + null)); + + testSchemaWithInvalidType( + createSchemaField( + "TripleUnionWithoutNull", + Schema.createUnion( + Schema.create(Schema.Type.BOOLEAN), + Schema.create(Schema.Type.INT), + Schema.create(Schema.Type.STRING)), + "", + null)); + + testSchemaWithInvalidType( + createSchemaField( + "DoubleUnionWithoutNull", + Schema.createUnion(Schema.create(Schema.Type.INT), Schema.create(Schema.Type.STRING)), + "", + null)); + + // Types that could eventually become supported... + + testSchemaWithInvalidType( + createSchemaField("StringArray", Schema.createArray(Schema.create(Schema.Type.STRING)), "", null)); + + testSchemaWithInvalidType( + createSchemaField("StringStringMap", Schema.createMap(Schema.create(Schema.Type.STRING)), "", null)); + + testSchemaWithInvalidType( + createSchemaField( + "Record", + Schema.createRecord("NestedRecord", "", "", false, Collections.emptyList()), + "", + null)); + } + + @Test + public void testUpsertStatement() { + Schema singleColumnKey = Schema.createRecord("MyKey", "", "", false, Collections.singletonList(getKey1())); + + List allFields = getAllValidFields(); + Schema schemaWithAllSupportedFieldTypes = Schema.createRecord("MyRecord", "", "", false, allFields); + + String upsertStatementForAllFields = AvroToSQL + .upsertStatement("MyRecord", singleColumnKey, schemaWithAllSupportedFieldTypes, Collections.emptySet()); + assertEquals(upsertStatementForAllFields, EXPECTED_UPSERT_STATEMENT_WITH_ALL_TYPES); + } + + private Schema.Field getKey1() { + return createSchemaField("key1", Schema.create(Schema.Type.INT), "", null); + } + + public static List getAllValidFields() { + List allFields = new ArrayList<>(); + + // Basic types + allFields.add(createSchemaField("fixedField", Schema.createFixed("MyFixed", "", "", 1), "", null)); + allFields.add(createSchemaField("stringField", Schema.create(Schema.Type.STRING), "", null)); + allFields.add(createSchemaField("bytesField", Schema.create(Schema.Type.BYTES), "", null)); + allFields.add(createSchemaField("intField", Schema.create(Schema.Type.INT), "", null)); + allFields.add(createSchemaField("longField", Schema.create(Schema.Type.LONG), "", null)); + allFields.add(createSchemaField("floatField", Schema.create(Schema.Type.FLOAT), "", null)); + allFields.add(createSchemaField("doubleField", Schema.create(Schema.Type.DOUBLE), "", null)); + allFields.add(createSchemaField("booleanField", Schema.create(Schema.Type.BOOLEAN), "", null)); + // allFields.add(createSchemaField("nullField", Schema.create(Schema.Type.NULL), "", null)); + + // Unions with null + List allOptionalFields = new ArrayList<>(); + for (Schema.Field field: allFields) { + if (field.schema().getType() == Schema.Type.NULL) { + // Madness? THIS -- IS -- SPARTAAAAAAAAAAAAAAAAAA!!!!!!!!! + continue; + } + + // Include both union branch orders + allOptionalFields.add( + createSchemaField( + field.name() + "Union1", + Schema.createUnion(Schema.create(Schema.Type.NULL), field.schema()), + "", + null)); + allOptionalFields.add( + createSchemaField( + field.name() + "Union2", + Schema.createUnion(field.schema(), Schema.create(Schema.Type.NULL)), + "", + null)); + } + allFields.addAll(allOptionalFields); + + return allFields; + } + + private void testSchemaWithInvalidType(Schema.Field invalidField) { + List allFields = getAllValidFields(); + allFields.add(invalidField); + + Schema.Field keyField1 = createSchemaField("key1", Schema.create(Schema.Type.INT), "", null); + Schema singleColumnKey = Schema.createRecord("MyKey", "", "", false, Collections.singletonList(keyField1)); + Schema valueSchema = Schema.createRecord("MyRecord", "", "", false, allFields); + + assertThrows( + IllegalArgumentException.class, + () -> SQLUtils.createTableStatement( + AvroToSQL + .getTableDefinition("MyRecord", singleColumnKey, valueSchema, Collections.emptySet(), FAIL, true))); + + String createTableStatement = SQLUtils.createTableStatement( + AvroToSQL.getTableDefinition("MyRecord", singleColumnKey, valueSchema, Collections.emptySet(), SKIP, false)); + assertEquals(createTableStatement, EXPECTED_CREATE_TABLE_STATEMENT_WITH_ALL_TYPES); + + String upsertStatement = + AvroToSQL.upsertStatement("MyRecord", singleColumnKey, valueSchema, Collections.emptySet()); + assertEquals(upsertStatement, EXPECTED_UPSERT_STATEMENT_WITH_ALL_TYPES); + } +} diff --git a/integrations/venice-duckdb/src/test/java/com/linkedin/venice/sql/SQLHelloWorldTest.java b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/sql/SQLHelloWorldTest.java new file mode 100644 index 00000000000..941e8990808 --- /dev/null +++ b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/sql/SQLHelloWorldTest.java @@ -0,0 +1,288 @@ +package com.linkedin.venice.sql; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.PreparedStatement; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + + +/** + * The aim of this class is just to test SQL engines, without any Venice-ism involved. + */ +public abstract class SQLHelloWorldTest { + protected abstract Connection getConnection() throws SQLException; + + protected abstract String getConnectionStringToPersistentDB(); + + /** + * Adapted from: https://duckdb.org/docs/api/java.html#querying + */ + @Test + public void test() throws SQLException { + try (Connection connection = getConnection(); Statement stmt = connection.createStatement()) { + // create a table + stmt.execute(createTableStatement("items", false)); + // insert two items into the table + stmt.execute(insertDataset1Statement("items")); + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM items")) { + assertValidityOfResultSet1(rs); + } + } + } + + /** + * This test verifies that DuckDB supports the same table swap technique as other RDBMS: + * + * BEGIN TRANSACTION; + * UPDATE TABLE current_version RENAME TO backup_version; + * UPDATE TABLE future_version RENAME TO current_version; + * COMMIT; + * + * This can be used as the basis for Venice version swaps. + */ + @Test + public void testVersionSwapViaTableRename() throws SQLException { + try (Connection connection = getConnection(); Statement stmt = connection.createStatement()) { + // create the current_version table + stmt.execute(createTableStatement("current_version", false)); + // insert two items into the table + stmt.execute(insertDataset1Statement("current_version")); + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM current_version")) { + assertValidityOfResultSet1(rs); + } + + // create the future_version table + stmt.execute(createTableStatement("future_version", false)); + // insert two items into the table + stmt.execute(insertDataset2Statement("future_version", false)); + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM future_version")) { + assertValidityOfResultSet2(rs); + } + + stmt.execute("BEGIN TRANSACTION;"); + stmt.execute("ALTER TABLE current_version RENAME to backup_version;"); + stmt.execute("ALTER TABLE future_version RENAME to current_version;"); + stmt.execute("COMMIT;"); + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM current_version")) { + assertValidityOfResultSet2(rs); + } + } + } + + /** + * This test verifies that DuckDB supports using view alteration as the mechanism for version swaps. + * + * This can be used as the basis for Venice version swaps. + */ + @Test + public void testVersionSwapViaViewAlteration() throws SQLException { + try (Connection connection = getConnection(); Statement stmt = connection.createStatement()) { + // create the current_version table + stmt.execute(createTableStatement("my_table_v1", false)); + // insert two items into the table + stmt.execute(insertDataset1Statement("my_table_v1")); + // create current_version view + stmt.execute(createViewStatement("my_table_v1")); + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM my_table_v1")) { + assertValidityOfResultSet1(rs); + } + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM my_table_current_version")) { + assertValidityOfResultSet1(rs); + } + + // create the future_version table + stmt.execute(createTableStatement("my_table_v2", false)); + // insert two items into the table + stmt.execute(insertDataset2Statement("my_table_v2", false)); + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM my_table_v2")) { + assertValidityOfResultSet2(rs); + } + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM my_table_current_version")) { + // The content of the view should remain unchanged as we have not swapped yet. + assertValidityOfResultSet1(rs); + } + + // SWAP! + stmt.execute(createViewStatement("my_table_v2")); + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM my_table_current_version")) { + assertValidityOfResultSet2(rs); + } + } + } + + @Test + public void testPrimaryKey() throws SQLException { + try (Connection connection = getConnection(); Statement stmt = connection.createStatement()) { + // create a table + stmt.execute(createTableStatement("items", true)); + // insert two items into the table + stmt.execute(insertDataset1Statement("items")); + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM items")) { + assertValidityOfResultSet1(rs); + } + + assertThrows(SQLException.class, () -> stmt.execute(insertDataset2Statement("items", false))); + } + } + + @Test + public void testUpsertStatement() throws SQLException { + try (Connection connection = getConnection(); Statement stmt = connection.createStatement()) { + // create a table + stmt.execute(createTableStatement("items", true)); + // insert two items into the table + stmt.execute(insertDataset1Statement("items")); + + stmt.execute(insertDataset2Statement("items", true)); + try (ResultSet rs = stmt.executeQuery("SELECT * FROM items")) { + assertValidityOfResultSet1WithUpsertDataSet2(rs); + } + } + } + + @DataProvider + public Object[][] upsertFlavors() { + String createTable1 = createTableStatement("items", true); + String createTable2 = "CREATE TABLE items (item VARCHAR, value DECIMAL(10, 2), count INTEGER, PRIMARY KEY (item))"; + String createTable3 = + "CREATE TABLE items (item VARCHAR, value DECIMAL(10, 2), count INTEGER, PRIMARY KEY (item, value))"; + String upsert1 = "INSERT OR REPLACE INTO items VALUES (?, ?, ?)"; + String upsert2 = + "INSERT INTO items VALUES (?, ?, ?) ON CONFLICT DO UPDATE SET value = EXCLUDED.value, count = EXCLUDED.count"; + String upsert3 = "INSERT INTO items VALUES (?, ?, ?) ON CONFLICT DO UPDATE SET count = EXCLUDED.count"; + return new Object[][] { { createTable1, upsert1 }, // + { createTable1, upsert2 }, // + { createTable2, upsert1 }, // + { createTable2, upsert2 }, // + { createTable3, upsert1 }, // + { createTable3, upsert3 } }; + } + + @Test(dataProvider = "upsertFlavors") + public void testUpsertPreparedStatement(String createTable, String upsert) throws SQLException { + try (Connection connection = getConnection(); Statement stmt = connection.createStatement()) { + // create a table + stmt.execute(createTable); + // insert two items into the table + stmt.execute(insertDataset1Statement("items")); + + try (PreparedStatement preparedStatement = connection.prepareStatement(upsert)) { + preparedStatement.setString(1, "jeans"); + preparedStatement.setDouble(2, 20.0); + preparedStatement.setInt(3, 2); + preparedStatement.execute(); + + preparedStatement.setString(1, "t-shirt"); + preparedStatement.setDouble(2, 42.2); + preparedStatement.setInt(3, 1); + preparedStatement.execute(); + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM items")) { + assertValidityOfResultSet1WithUpsertDataSet2(rs); + } + } + } + } + + @Test + public void testPersistence() throws SQLException { + String connectionString = getConnectionStringToPersistentDB(); + System.out.println(connectionString); + try (Connection connection = DriverManager.getConnection(connectionString); + Statement stmt = connection.createStatement()) { + // create a table + stmt.execute(createTableStatement("items", true)); + // insert two items into the table + stmt.execute(insertDataset1Statement("items")); + + try (ResultSet rs = stmt.executeQuery("SELECT * FROM items")) { + assertValidityOfResultSet1(rs); + } + } + + try (Connection connection = DriverManager.getConnection(connectionString); + Statement stmt = connection.createStatement()) { + // Table and data should still be there even though we closed the DuckDBConnection + try (ResultSet rs = stmt.executeQuery("SELECT * FROM items")) { + assertValidityOfResultSet1(rs); + } + } + } + + protected String createTableStatement(String tableName, boolean primaryKey) { + String pk = primaryKey ? " PRIMARY KEY" : ""; + return "CREATE TABLE " + tableName + " (item VARCHAR" + pk + ", value DECIMAL(10, 2), count INTEGER)"; + } + + private String createViewStatement(String tableName) { + return "CREATE OR REPLACE VIEW my_table_current_version AS SELECT * FROM " + tableName + ";"; + } + + private String insertDataset1Statement(String tableName) { + return "INSERT INTO " + tableName + " VALUES ('jeans', 20.0, 1), ('hammer', 42.2, 2)"; + } + + private String insertDataset2Statement(String tableName, boolean upsert) { + String orReplace = upsert ? " OR REPLACE" : ""; + return "INSERT" + orReplace + " INTO " + tableName + " VALUES ('jeans', 20.0, 2), ('t-shirt', 42.2, 1)"; + } + + protected void assertValidityOfResultSet1(ResultSet rs) throws SQLException { + assertTrue(rs.next(), "There should be a first row!"); + assertEquals(rs.getString(1), "jeans"); + assertEquals(rs.getInt(3), 1); + + assertTrue(rs.next(), "There should be a second row!"); + assertEquals(rs.getString(1), "hammer"); + assertEquals(rs.getInt(3), 2); + + assertFalse(rs.next(), "There should only be two rows!"); + } + + private void assertValidityOfResultSet2(ResultSet rs) throws SQLException { + assertTrue(rs.next(), "There should be a first row!"); + assertEquals(rs.getString(1), "jeans"); + assertEquals(rs.getInt(3), 2); + + assertTrue(rs.next(), "There should be a second row!"); + assertEquals(rs.getString(1), "t-shirt"); + assertEquals(rs.getInt(3), 1); + + assertFalse(rs.next(), "There should only be two rows!"); + } + + private void assertValidityOfResultSet1WithUpsertDataSet2(ResultSet rs) throws SQLException { + assertTrue(rs.next(), "There should be a first row!"); + assertEquals(rs.getString(1), "jeans"); + assertEquals(rs.getInt(3), 2); + + assertTrue(rs.next(), "There should be a second row!"); + assertEquals(rs.getString(1), "hammer"); + assertEquals(rs.getInt(3), 2); + + assertTrue(rs.next(), "There should be a third row!"); + assertEquals(rs.getString(1), "t-shirt"); + assertEquals(rs.getInt(3), 1); + + assertFalse(rs.next(), "There should only be three rows!"); + } +} diff --git a/integrations/venice-duckdb/src/test/java/com/linkedin/venice/sql/SQLUtilsTest.java b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/sql/SQLUtilsTest.java new file mode 100644 index 00000000000..d196f014110 --- /dev/null +++ b/integrations/venice-duckdb/src/test/java/com/linkedin/venice/sql/SQLUtilsTest.java @@ -0,0 +1,183 @@ +package com.linkedin.venice.sql; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNotSame; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; + +import java.sql.Connection; +import java.sql.JDBCType; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import org.testng.annotations.Test; + + +public abstract class SQLUtilsTest { + private static final String TABLE_NAME = "items"; + private static final String KEY1_NAME = "key1"; + private static final String KEY2_NAME = "key2"; + private static final String VAL1_NAME = "val1"; + private static final String VAL2_NAME = "val2"; + + protected abstract Connection getConnection() throws SQLException; + + @Test + public void pojoTesting() { + // Basics + TableDefinition tableDefinition = getTableDefinition(); + assertEquals(tableDefinition.getName(), TABLE_NAME); + assertEquals(tableDefinition.getColumns().size(), 4); + assertEquals(tableDefinition.getPrimaryKeyColumns().size(), 2); + + int expectedJdbcIndex = 1; + ColumnDefinition key1 = tableDefinition.getColumnByJdbcIndex(expectedJdbcIndex); + assertNotNull(key1); + assertEquals(key1.getName(), KEY1_NAME); + assertEquals(key1.getType(), JDBCType.INTEGER); + assertFalse(key1.isNullable()); + assertEquals(key1.getIndexType(), IndexType.PRIMARY_KEY); + assertNull(key1.getDefaultValue()); + assertNull(key1.getExtra()); + assertEquals(key1.getJdbcIndex(), expectedJdbcIndex++); + + ColumnDefinition key2 = tableDefinition.getColumnByJdbcIndex(expectedJdbcIndex); + assertNotNull(key2); + assertEquals(key2.getName(), KEY2_NAME); + assertEquals(key2.getType(), JDBCType.BIGINT); + assertFalse(key2.isNullable()); + assertEquals(key2.getIndexType(), IndexType.PRIMARY_KEY); + assertNull(key2.getDefaultValue()); + assertNull(key2.getExtra()); + assertEquals(key2.getJdbcIndex(), expectedJdbcIndex++); + + ColumnDefinition val1 = tableDefinition.getColumnByJdbcIndex(expectedJdbcIndex); + assertNotNull(val1); + assertEquals(val1.getName(), VAL1_NAME); + assertEquals(val1.getType(), JDBCType.VARCHAR); + assertTrue(val1.isNullable()); + assertNull(val1.getIndexType()); + assertNull(val1.getDefaultValue()); + assertNull(val1.getExtra()); + assertEquals(val1.getJdbcIndex(), expectedJdbcIndex++); + + ColumnDefinition val2 = tableDefinition.getColumnByJdbcIndex(expectedJdbcIndex); + assertNotNull(val2); + assertEquals(val2.getName(), VAL2_NAME); + assertEquals(val2.getType(), JDBCType.DOUBLE); + assertTrue(val2.isNullable()); + assertNull(val2.getIndexType()); + assertNull(val2.getDefaultValue()); + assertNull(val2.getExtra()); + assertEquals(val2.getJdbcIndex(), expectedJdbcIndex++); + + // Equality checks + TableDefinition otherTableDefinition = getTableDefinition(); + assertNotSame(tableDefinition, otherTableDefinition); + assertEquals(tableDefinition, otherTableDefinition); + assertEquals(tableDefinition, tableDefinition); + + TableDefinition differentTableDefinition = new TableDefinition(TABLE_NAME, Collections.singletonList(key1)); + assertNotEquals(tableDefinition, differentTableDefinition); + + assertNotEquals(tableDefinition, null); + assertNotEquals(null, tableDefinition); + assertNotEquals(tableDefinition, ""); + assertNotEquals("", tableDefinition); + + ColumnDefinition otherKey1 = otherTableDefinition.getColumnByName(KEY1_NAME); + assertNotSame(key1, otherKey1); + assertEquals(key1, otherKey1); + assertEquals(key1, key1); + + ColumnDefinition otherKey2 = otherTableDefinition.getColumnByName(KEY2_NAME); + assertNotSame(key2, otherKey2); + assertEquals(key2, otherKey2); + + ColumnDefinition otherVal1 = otherTableDefinition.getColumnByName(VAL1_NAME); + assertNotSame(val1, otherVal1); + assertEquals(val1, otherVal1); + + ColumnDefinition otherVal2 = otherTableDefinition.getColumnByName(VAL2_NAME); + assertNotSame(val2, otherVal2); + assertEquals(val2, otherVal2); + + assertNull(otherTableDefinition.getColumnByName("bogus")); + + assertNotEquals(key1, null); + assertNotEquals(null, key1); + assertNotEquals(key1, ""); + assertNotEquals("", key1); + + // Invalid inputs + assertThrows(NullPointerException.class, () -> new TableDefinition(null, Collections.emptyList())); + assertThrows(IllegalArgumentException.class, () -> new TableDefinition(TABLE_NAME, Collections.emptyList())); + + assertThrows(IllegalArgumentException.class, () -> new ColumnDefinition(KEY1_NAME, JDBCType.INTEGER, 0)); + assertThrows(IllegalArgumentException.class, () -> new ColumnDefinition(VAL1_NAME, JDBCType.VARCHAR, 0)); + + List outOfOrderColumns = new ArrayList<>(); + outOfOrderColumns.add(key2); + outOfOrderColumns.add(key1); + assertThrows(IllegalArgumentException.class, () -> new TableDefinition(TABLE_NAME, outOfOrderColumns)); + + List sparseColumns = new ArrayList<>(); + sparseColumns.add(key1); + sparseColumns.add(key2); + sparseColumns.add(1, null); + assertThrows(IllegalArgumentException.class, () -> new TableDefinition(TABLE_NAME, sparseColumns)); + + assertThrows(IndexOutOfBoundsException.class, () -> tableDefinition.getColumnByJdbcIndex(0)); + assertThrows(IndexOutOfBoundsException.class, () -> tableDefinition.getColumnByJdbcIndex(5)); + } + + @Test + public void testGetTableDefinition() throws SQLException { + try (Connection connection = getConnection(); Statement statement = connection.createStatement()) { + // Starting point: empty + assertNull(SQLUtils.getTableDefinition(TABLE_NAME, connection)); + + String differentTable = "not_the_table_we_want"; + statement.execute( + "CREATE TABLE " + differentTable + " (" + KEY1_NAME + " INTEGER, " + KEY2_NAME + " BIGINT, " + VAL1_NAME + + " VARCHAR UNIQUE, " + VAL2_NAME + " DOUBLE, PRIMARY KEY (key1, key2));"); + + assertNull(SQLUtils.getTableDefinition(TABLE_NAME, connection)); + + statement.execute( + "CREATE TABLE " + TABLE_NAME + " (" + KEY1_NAME + " INTEGER, " + KEY2_NAME + " BIGINT, " + VAL1_NAME + + " VARCHAR, " + VAL2_NAME + " DOUBLE, PRIMARY KEY (key1, key2));"); + + TableDefinition tableDefinition = SQLUtils.getTableDefinition(TABLE_NAME, connection); + assertNotNull(tableDefinition); + + TableDefinition expectedTableDefinition = getTableDefinition(); + assertEquals(tableDefinition, expectedTableDefinition); + + TableDefinition differentTableDefinition = SQLUtils.getTableDefinition(differentTable, connection); + assertNotNull(differentTableDefinition); + assertNotEquals(tableDefinition, differentTableDefinition); + } + } + + private TableDefinition getTableDefinition() { + int jdbcIndex = 1; + ColumnDefinition key1 = + new ColumnDefinition(KEY1_NAME, JDBCType.INTEGER, false, IndexType.PRIMARY_KEY, jdbcIndex++); + ColumnDefinition key2 = new ColumnDefinition(KEY2_NAME, JDBCType.BIGINT, false, IndexType.PRIMARY_KEY, jdbcIndex++); + ColumnDefinition val1 = new ColumnDefinition(VAL1_NAME, JDBCType.VARCHAR, jdbcIndex++); + ColumnDefinition val2 = new ColumnDefinition(VAL2_NAME, JDBCType.DOUBLE, jdbcIndex++); + List columns = new ArrayList<>(); + columns.add(key1); + columns.add(key2); + columns.add(val1); + columns.add(val2); + return new TableDefinition(TABLE_NAME, columns); + } +} diff --git a/integrations/venice-samza/src/test/java/com/linkedin/venice/samza/VeniceSystemProducerTest.java b/integrations/venice-samza/src/test/java/com/linkedin/venice/samza/VeniceSystemProducerTest.java index 1af3cb0588c..d309b527c40 100644 --- a/integrations/venice-samza/src/test/java/com/linkedin/venice/samza/VeniceSystemProducerTest.java +++ b/integrations/venice-samza/src/test/java/com/linkedin/venice/samza/VeniceSystemProducerTest.java @@ -2,7 +2,16 @@ import static com.linkedin.venice.ConfigKeys.KAFKA_BOOTSTRAP_SERVERS; import static com.linkedin.venice.pubsub.adapter.kafka.producer.ApacheKafkaProducerConfig.KAFKA_BUFFER_MEMORY; -import static org.mockito.Mockito.*; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.anyBoolean; +import static org.mockito.Mockito.anyLong; +import static org.mockito.Mockito.anyString; +import static org.mockito.Mockito.doAnswer; +import static org.mockito.Mockito.doNothing; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.spy; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertNull; diff --git a/internal/alpini/common/alpini-common-base/src/test/java/com/linkedin/alpini/base/misc/TestTimeScheduledThreadPoolExecutor.java b/internal/alpini/common/alpini-common-base/src/test/java/com/linkedin/alpini/base/misc/TestTimeScheduledThreadPoolExecutor.java index 0db38ded824..41f5665e0da 100644 --- a/internal/alpini/common/alpini-common-base/src/test/java/com/linkedin/alpini/base/misc/TestTimeScheduledThreadPoolExecutor.java +++ b/internal/alpini/common/alpini-common-base/src/test/java/com/linkedin/alpini/base/misc/TestTimeScheduledThreadPoolExecutor.java @@ -26,7 +26,9 @@ */ @Test(singleThreaded = true) public class TestTimeScheduledThreadPoolExecutor { - @Test(groups = "unit") + private static final long DEFAULT_TIMEOUT = 90_000L; + + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testQueueDrain() throws Exception { BlockingQueue queue = new TimeScheduledThreadPoolExecutor.DelayedWorkQueue(); try { @@ -83,7 +85,7 @@ class Task extends TimeScheduledThreadPoolExecutor.AbstractFutureTask { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testQueueTimeSkip() throws Exception { BlockingQueue queue = new TimeScheduledThreadPoolExecutor.DelayedWorkQueue(); try { @@ -117,7 +119,7 @@ class Task extends TimeScheduledThreadPoolExecutor.AbstractFutureTask { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testQueueLongTimeSkip() throws Exception { BlockingQueue queue = new TimeScheduledThreadPoolExecutor.DelayedWorkQueue(); AtomicBoolean shutdown = new AtomicBoolean(); @@ -166,7 +168,7 @@ class Task extends TimeScheduledThreadPoolExecutor.AbstractFutureTask { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testShutdown() throws InterruptedException { TimeScheduledThreadPoolExecutor executor = new TimeScheduledThreadPoolExecutor(1); Runnable mock = Mockito.mock(Runnable.class); @@ -217,7 +219,7 @@ public void testShutdown() throws InterruptedException { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testBasicTimeSkip() throws InterruptedException { ScheduledExecutorService executor = new TimeScheduledThreadPoolExecutor(1); try { @@ -249,7 +251,7 @@ public void testBasicTimeSkip() throws InterruptedException { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testLongTimeSkip() throws InterruptedException { ScheduledExecutorService executor = new TimeScheduledThreadPoolExecutor(1); Thread timeWarp = new Thread(() -> { @@ -362,7 +364,7 @@ private void scheduleAtTheEndOfTime(TimeScheduledThreadPoolExecutor pool, Runnab } } - @Test(groups = "unit", dataProvider = "testCombo1") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, dataProvider = "testCombo1") public void testScheduleTimeOverflow1(int nowHow, int thenHow) throws InterruptedException { final TimeScheduledThreadPoolExecutor pool = new TimeScheduledThreadPoolExecutor(1); try { @@ -397,7 +399,7 @@ public void testScheduleTimeOverflow1(int nowHow, int thenHow) throws Interrupte } } - @Test(groups = "unit", dataProvider = "testCombo2") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, dataProvider = "testCombo2") public void testScheduleTimeOverflow2(int nowHow) throws InterruptedException { final int nowHowCopy = nowHow; final TimeScheduledThreadPoolExecutor pool = new TimeScheduledThreadPoolExecutor(1); diff --git a/internal/alpini/common/alpini-common-base/src/test/java/com/linkedin/alpini/base/registry/TestResourceRegistry.java b/internal/alpini/common/alpini-common-base/src/test/java/com/linkedin/alpini/base/registry/TestResourceRegistry.java index bfac41ead2b..1882a446ad5 100644 --- a/internal/alpini/common/alpini-common-base/src/test/java/com/linkedin/alpini/base/registry/TestResourceRegistry.java +++ b/internal/alpini/common/alpini-common-base/src/test/java/com/linkedin/alpini/base/registry/TestResourceRegistry.java @@ -30,6 +30,8 @@ * @author Antony T Curtis <acurtis@linkedin.com> */ public class TestResourceRegistry { + private static final long DEFAULT_TIMEOUT = 90_000L; + public static interface MockResource extends ShutdownableResource { } @@ -153,34 +155,34 @@ public void beforeMethod() { Mockito.reset(factory1); } - @Test(groups = "unit", expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "Factory missing generic type") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "Factory missing generic type") public void testRegisterBadFactory() { Assert.assertNull(ResourceRegistry.registerFactory(BadFactory.class, Mockito.mock(BadFactory.class))); } - @Test(groups = "unit", expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = ".*BadClass is not an interface.*") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = ".*BadClass is not an interface.*") public void testRegisterBadClass() { Assert.assertNull(ResourceRegistry.registerFactory(BadClass.class, new BadClass())); } - @Test(groups = "unit", expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "Factory should be an instance of the Factory class.") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = "Factory should be an instance of the Factory class.") @SuppressWarnings("unchecked") public void testRegisterBadMismatchFactory() { Assert.assertNull( ResourceRegistry.registerFactory((Class) EmptyFactory.class, (ResourceRegistry.Factory) new BadClass())); } - @Test(groups = "unit", expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = ".*EmptyFactory does not declare any methods") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = ".*EmptyFactory does not declare any methods") public void testRegisterEmptyFactory() { Assert.assertNull(ResourceRegistry.registerFactory(EmptyFactory.class, Mockito.mock(EmptyFactory.class))); } - @Test(groups = "unit", expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = ".*InvalidFactory does not return a \\[.*MockResource\\] class") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, expectedExceptions = IllegalArgumentException.class, expectedExceptionsMessageRegExp = ".*InvalidFactory does not return a \\[.*MockResource\\] class") public void testRegisterFactory() { Assert.assertNull(ResourceRegistry.registerFactory(InvalidFactory.class, Mockito.mock(InvalidFactory.class))); } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testRegisterRace() throws InterruptedException { final ResourceRegistry reg = new ResourceRegistry(); try { @@ -226,28 +228,28 @@ public void run() { } } - @Test(groups = "unit", expectedExceptions = Error.class, expectedExceptionsMessageRegExp = "Bad foo always happens") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, expectedExceptions = Error.class, expectedExceptionsMessageRegExp = "Bad foo always happens") public void testFactoryStaticError() { ResourceRegistry reg = new ResourceRegistry(); FactoryStaticError factory = reg.factory(FactoryStaticError.class); Assert.fail("Should not get here: " + factory); } - @Test(groups = "unit", expectedExceptions = NumberFormatException.class) + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, expectedExceptions = NumberFormatException.class) public void testFactoryStaticRuntimeException() { ResourceRegistry reg = new ResourceRegistry(); FactoryStaticRuntimeException factory = reg.factory(FactoryStaticRuntimeException.class); Assert.fail("Should not get here: " + factory); } - @Test(groups = "unit", expectedExceptions = NullPointerException.class) + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, expectedExceptions = NullPointerException.class) public void testFactoryNull() { ResourceRegistry reg = new ResourceRegistry(); ResourceRegistry.Factory factory = reg.factory(null); Assert.fail("Should not get here: " + factory); } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testRegisterMockitoFactory1() throws InterruptedException, TimeoutException { Assert.assertSame(ResourceRegistry.registerFactory(MockFactory1.class, factory1), factory1); MockResource mockResource = Mockito.mock(MockResource.class); @@ -321,7 +323,7 @@ public Object answer(InvocationOnMock invocation) throws Throwable { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testRegisterMockitoFactory5() throws InterruptedException, TimeoutException { Assert.assertSame(ResourceRegistry.registerFactory(MockFactory5.class, factory5), factory5); MockResource2 mockResource2 = Mockito.mock(MockResource2.class); @@ -379,7 +381,7 @@ public void testRegisterMockitoFactory5() throws InterruptedException, TimeoutEx } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testRegisterMockitoFactory1BadWaitForShutdown() throws InterruptedException, TimeoutException { Assert.assertSame(ResourceRegistry.registerFactory(MockFactory1.class, factory1), factory1); MockResource mockResource = Mockito.mock(MockResource.class); @@ -433,7 +435,7 @@ public Boolean answer(InvocationOnMock invocation) throws Throwable { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testRegisterMockitoFactory1BadShutdown() throws InterruptedException, TimeoutException { Assert.assertSame(ResourceRegistry.registerFactory(MockFactory1.class, factory1), factory1); MockResource mockResource = Mockito.mock(MockResource.class); @@ -478,7 +480,7 @@ public Object answer(InvocationOnMock invocation) throws Throwable { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testRegisterMockitoFactory1Dormouse() throws InterruptedException, TimeoutException { Assert.assertSame(ResourceRegistry.registerFactory(MockFactory1.class, factory1), factory1); MockResource mockResource = Mockito.mock(MockResource.class); @@ -523,7 +525,7 @@ public Object answer(InvocationOnMock invocation) throws Throwable { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testRegisterMockitoFactory1Timeout() throws InterruptedException, TimeoutException { Assert.assertSame(ResourceRegistry.registerFactory(MockFactory1.class, factory1), factory1); MockResource mockResource = Mockito.mock(MockResource.class); @@ -614,7 +616,7 @@ public Object answer(InvocationOnMock invocation) throws Throwable { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testRemoveShutdownableResource() throws InterruptedException { ResourceRegistry reg = new ResourceRegistry(); ShutdownableResource[] res = new ShutdownableResource[5]; @@ -673,7 +675,7 @@ public Boolean answer(InvocationOnMock invocation) throws Throwable { Assert.assertEquals(latch[4].getCount(), 1); } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testRemoveShutdownable() throws InterruptedException { ResourceRegistry reg = new ResourceRegistry(); Shutdownable[] res = new Shutdownable[4]; @@ -708,7 +710,7 @@ public Object answer(InvocationOnMock invocation) throws Throwable { Assert.assertEquals(latch[3].getCount(), 0); } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testRemoveSyncShutdownable() throws InterruptedException { ResourceRegistry reg = new ResourceRegistry(); SyncShutdownable[] res = new SyncShutdownable[4]; @@ -743,7 +745,7 @@ public Object answer(InvocationOnMock invocation) throws Throwable { Assert.assertEquals(latch[3].getCount(), 0); } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testSortedShutdownCombinations() throws InterruptedException, TimeoutException { Assert.assertSame(ResourceRegistry.registerFactory(MockFactory1.class, factory1), factory1); MockResource[] mockResources = { Mockito.mock(MockResourceFirst.class), Mockito.mock(MockResource.class), @@ -871,7 +873,7 @@ public Object answer(InvocationOnMock invocation) throws Throwable { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testReRegisterInvalidFactory() { ReregisterFactory factory = Mockito.mock(ReregisterFactory.class); Assert.assertSame(ResourceRegistry.registerFactory(ReregisterFactory.class, factory), factory); @@ -881,7 +883,7 @@ public void testReRegisterInvalidFactory() { Assert.assertSame(ResourceRegistry.registerFactory(ReregisterFactory.class, factory), factory); } - @Test(groups = "unit", expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = "Factory not registered for interface.*") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, expectedExceptions = IllegalStateException.class, expectedExceptionsMessageRegExp = "Factory not registered for interface.*") public void testRunNotRegistered() { ResourceRegistry reg = new ResourceRegistry(false); try { @@ -891,7 +893,7 @@ public void testRunNotRegistered() { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testShutdownQuick() throws TimeoutException, InterruptedException { ResourceRegistry reg = new ResourceRegistry(false); Assert.assertFalse(reg.isShutdown()); @@ -903,7 +905,7 @@ public void testShutdownQuick() throws TimeoutException, InterruptedException { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testToString() { ResourceRegistry reg = new ResourceRegistry(); String str; @@ -913,7 +915,7 @@ public void testToString() { reg.shutdown(); } - @Test(groups = "unit", expectedExceptions = IllegalStateException.class) + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, expectedExceptions = IllegalStateException.class) public void testPrematureWaitForShutdown1() throws InterruptedException { ResourceRegistry reg = new ResourceRegistry(); try { @@ -923,7 +925,7 @@ public void testPrematureWaitForShutdown1() throws InterruptedException { } } - @Test(groups = "unit", expectedExceptions = IllegalStateException.class) + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT, expectedExceptions = IllegalStateException.class) public void testPrematureWaitForShutdown2() throws InterruptedException, TimeoutException { ResourceRegistry reg = new ResourceRegistry(); try { @@ -933,7 +935,7 @@ public void testPrematureWaitForShutdown2() throws InterruptedException, Timeout } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testUnilateralGlobalShutdown() throws InterruptedException, TimeoutException { MockFactory2 factory = Mockito.mock(MockFactory2.class); Assert.assertSame(ResourceRegistry.registerFactory(MockFactory2.class, factory), factory); @@ -1048,7 +1050,7 @@ public Object answer(InvocationOnMock invocation) throws Throwable { } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testUnilateralGlobalShutdownInterruptable() throws InterruptedException, TimeoutException { MockFactory3 factory = Mockito.mock(MockFactory3.class); Assert.assertSame(ResourceRegistry.registerFactory(MockFactory3.class, factory), factory); @@ -1157,7 +1159,7 @@ public Void call() throws Exception { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testUnilateralGlobalShutdownDelay() throws InterruptedException, TimeoutException { MockFactory4 factory = Mockito.mock(MockFactory4.class); Assert.assertSame(ResourceRegistry.registerFactory(MockFactory4.class, factory), factory); @@ -1263,7 +1265,7 @@ public Object answer(InvocationOnMock invocation) throws Throwable { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testResourceGarbageCollection() throws InterruptedException, TimeoutException { try { MockFactoryGC factoryGC = Mockito.mock(MockFactoryGC.class); @@ -1332,7 +1334,7 @@ public Object answer(InvocationOnMock invocation) throws Throwable { } } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testWaitToStartShutdown() throws InterruptedException, TimeoutException, ExecutionException { final ResourceRegistry registry = new ResourceRegistry(); final RunnableFuture future = new FutureTask(new Callable() { @@ -1357,7 +1359,7 @@ public Object call() throws Exception { future.get(1, TimeUnit.MILLISECONDS); } - @Test(groups = "unit") + @Test(groups = "unit", timeOut = DEFAULT_TIMEOUT) public void testWaitToStartShutdownTimeout() throws InterruptedException, TimeoutException, ExecutionException { final ResourceRegistry registry = new ResourceRegistry(); final CountDownLatch startLatch = new CountDownLatch(2); diff --git a/internal/venice-client-common/src/main/java/com/linkedin/venice/stats/AbstractVeniceStats.java b/internal/venice-client-common/src/main/java/com/linkedin/venice/stats/AbstractVeniceStats.java index 9bbf2dee894..a1ffc7b92aa 100644 --- a/internal/venice-client-common/src/main/java/com/linkedin/venice/stats/AbstractVeniceStats.java +++ b/internal/venice-client-common/src/main/java/com/linkedin/venice/stats/AbstractVeniceStats.java @@ -136,7 +136,7 @@ protected Sensor registerSensor( } } } else { - String metricName = sensorFullName + "." + metricNameSuffix(stat); + String metricName = getMetricFullName(sensor, stat); if (metricsRepository.getMetric(metricName) == null) { sensor.add(metricName, stat, config); } @@ -147,6 +147,10 @@ protected Sensor registerSensor( }); } + protected String getMetricFullName(Sensor sensor, MeasurableStat stat) { + return sensor.name() + "." + metricNameSuffix(stat); + } + /** * N.B.: {@link LongAdderRateGauge} is just an implementation detail, and we do not wish to alter metric names * due to it, so we call it the same as {@link Rate}. Same for {@link AsyncGauge}, we don't want to alter any existing diff --git a/internal/venice-client-common/src/test/java/com/linkedin/venice/compute/ComputeUtilsTest.java b/internal/venice-client-common/src/test/java/com/linkedin/venice/compute/ComputeUtilsTest.java index c38c40b98b4..f9dabf6d6ea 100644 --- a/internal/venice-client-common/src/test/java/com/linkedin/venice/compute/ComputeUtilsTest.java +++ b/internal/venice-client-common/src/test/java/com/linkedin/venice/compute/ComputeUtilsTest.java @@ -1,7 +1,10 @@ package com.linkedin.venice.compute; -import static com.linkedin.venice.utils.TestWriteUtils.*; -import static org.testng.Assert.*; +import static com.linkedin.venice.utils.TestWriteUtils.loadFileAsString; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; import com.linkedin.avro.api.PrimitiveFloatList; import com.linkedin.avro.fastserde.primitive.PrimitiveFloatArrayList; diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/ConfigKeys.java b/internal/venice-common/src/main/java/com/linkedin/venice/ConfigKeys.java index 02f0d293443..eab209a0e76 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/ConfigKeys.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/ConfigKeys.java @@ -225,6 +225,27 @@ private ConfigKeys() { // What tags to assign to a controller instance public static final String CONTROLLER_INSTANCE_TAG_LIST = "controller.instance.tag.list"; + /** + * Whether to enable gRPC server in controller or not. + */ + public static final String CONTROLLER_GRPC_SERVER_ENABLED = "controller.grpc.server.enabled"; + + /** + * A port for the controller to listen on for incoming requests. On this port, the controller will + * server non-ssl requests. + */ + public static final String CONTROLLER_ADMIN_GRPC_PORT = "controller.admin.grpc.port"; + /** + * A port for the controller to listen on for incoming requests. On this port, the controller will + * only serve ssl requests. + */ + public static final String CONTROLLER_ADMIN_SECURE_GRPC_PORT = "controller.admin.secure.grpc.port"; + + /** + * Number of threads to use for the gRPC server in controller. + */ + public static final String CONTROLLER_GRPC_SERVER_THREAD_COUNT = "controller.grpc.server.thread.count"; + /** List of forbidden admin paths */ public static final String CONTROLLER_DISABLED_ROUTES = "controller.cluster.disabled.routes"; @@ -759,6 +780,13 @@ private ConfigKeys() { public static final String SERVER_DB_READ_ONLY_FOR_BATCH_ONLY_STORE_ENABLED = "server.db.read.only.for.batch.only.store.enabled"; public static final String SERVER_RESET_ERROR_REPLICA_ENABLED = "server.reset.error.replica.enabled"; + + public static final String SERVER_ADAPTIVE_THROTTLER_ENABLED = "server.adaptive.throttler.enabled"; + public static final String SERVER_ADAPTIVE_THROTTLER_SIGNAL_IDLE_THRESHOLD = + "server.adaptive.throttler.signal.idle.threshold"; + public static final String SERVER_ADAPTIVE_THROTTLER_SINGLE_GET_LATENCY_THRESHOLD = + "server.adaptive.throttler.single.get.latency.threshold"; + /** * A list of fully-qualified class names of all stats classes that needs to be initialized in isolated ingestion process, * separated by comma. This config will help isolated ingestion process to register extra stats needed for monitoring, @@ -2351,4 +2379,7 @@ private ConfigKeys() { public static final String SERVER_DELETE_UNASSIGNED_PARTITIONS_ON_STARTUP = "server.delete.unassigned.partitions.on.startup"; + + public static final String CONTROLLER_ENABLE_HYBRID_STORE_PARTITION_COUNT_UPDATE = + "controller.enable.hybrid.store.partition.count.update"; } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/acl/NoOpDynamicAccessController.java b/internal/venice-common/src/main/java/com/linkedin/venice/acl/NoOpDynamicAccessController.java new file mode 100644 index 00000000000..794ec55e41a --- /dev/null +++ b/internal/venice-common/src/main/java/com/linkedin/venice/acl/NoOpDynamicAccessController.java @@ -0,0 +1,76 @@ +package com.linkedin.venice.acl; + +import java.security.cert.X509Certificate; +import java.util.Collections; +import java.util.List; +import java.util.Set; + + +/** + * A no-op implementation of {@link DynamicAccessController}. + */ +public class NoOpDynamicAccessController implements DynamicAccessController { + public static final String USER_UNKNOWN = "USER_UNKNOWN"; + + public static final NoOpDynamicAccessController INSTANCE = new NoOpDynamicAccessController(); + + private NoOpDynamicAccessController() { + } + + @Override + public DynamicAccessController init(List resources) { + return this; + } + + @Override + public boolean hasAccess(X509Certificate clientCert, String resource, String method) throws AclException { + return true; + } + + @Override + public boolean hasAccessToTopic(X509Certificate clientCert, String resource, String method) throws AclException { + return true; + } + + @Override + public boolean hasAccessToAdminOperation(X509Certificate clientCert, String operation) throws AclException { + return true; + } + + @Override + public boolean isAllowlistUsers(X509Certificate clientCert, String resource, String method) { + return true; + } + + @Override + public String getPrincipalId(X509Certificate clientCert) { + if (clientCert != null && clientCert.getSubjectX500Principal() != null) { + return clientCert.getSubjectX500Principal().getName(); + } + return USER_UNKNOWN; + } + + @Override + public boolean hasAcl(String resource) throws AclException { + return true; + } + + @Override + public void addAcl(String resource) throws AclException { + } + + @Override + public void removeAcl(String resource) throws AclException { + + } + + @Override + public Set getAccessControlledResources() { + return Collections.emptySet(); + } + + @Override + public boolean isFailOpen() { + return false; + } +} diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/common/VeniceSystemStoreType.java b/internal/venice-common/src/main/java/com/linkedin/venice/common/VeniceSystemStoreType.java index 38a3dfd04a3..756a7b08059 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/common/VeniceSystemStoreType.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/common/VeniceSystemStoreType.java @@ -6,6 +6,7 @@ import com.linkedin.venice.authorization.Permission; import com.linkedin.venice.authorization.Resource; import com.linkedin.venice.meta.Store; +import com.linkedin.venice.meta.StoreName; import com.linkedin.venice.pushstatus.PushStatusKey; import com.linkedin.venice.pushstatus.PushStatusValue; import com.linkedin.venice.serialization.avro.AvroProtocolDefinition; @@ -164,7 +165,7 @@ public Method getClientAccessMethod() { */ public AclBinding generateSystemStoreAclBinding(AclBinding regularStoreAclBinding) { String regularStoreName = regularStoreAclBinding.getResource().getName(); - if (!Store.isValidStoreName(regularStoreName)) { + if (!StoreName.isValidStoreName(regularStoreName)) { throw new UnsupportedOperationException( "Cannot generate system store AclBinding for a non-store resource: " + regularStoreName); } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerApiConstants.java b/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerApiConstants.java index a755a73d9ed..5289dd7aa0d 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerApiConstants.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerApiConstants.java @@ -7,7 +7,11 @@ public class ControllerApiConstants { public static final String SOURCE_GRID_FABRIC = "source_grid_fabric"; public static final String BATCH_JOB_HEARTBEAT_ENABLED = "batch_job_heartbeat_enabled"; - public static final String NAME = "store_name"; + public static final String STORE_NAME = "store_name"; + /** + * @deprecated Use {@link #STORE_NAME} instead. + */ + public static final String NAME = STORE_NAME; public static final String STORE_PARTITION = "store_partition"; public static final String STORE_VERSION = "store_version"; public static final String OWNER = "owner"; diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerClient.java b/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerClient.java index 61ecaba454a..0cf322611a1 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerClient.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/ControllerClient.java @@ -201,7 +201,8 @@ protected String discoverLeaderController() { String leaderControllerUrl = transport.request(url, ControllerRoute.LEADER_CONTROLLER, newParams(), LeaderControllerResponse.class) .getUrl(); - LOGGER.info("Discovered leader controller {} from {}", leaderControllerUrl, url); + LOGGER + .info("Discovered leader controller: {} from: {} for cluster: {}", leaderControllerUrl, url, clusterName); return leaderControllerUrl; } catch (Exception e) { LOGGER.warn("Unable to discover leader controller from {}", url); diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/LeaderControllerResponse.java b/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/LeaderControllerResponse.java index 7332f1581b3..443025578eb 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/LeaderControllerResponse.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/LeaderControllerResponse.java @@ -5,6 +5,8 @@ public class LeaderControllerResponse private String cluster; private String url; private String secureUrl = null; + private String grpcUrl = null; + private String secureGrpcUrl = null; public String getCluster() { return cluster; @@ -29,4 +31,20 @@ public String getSecureUrl() { public void setSecureUrl(String url) { this.secureUrl = url; } + + public void setGrpcUrl(String url) { + this.grpcUrl = url; + } + + public String getGrpcUrl() { + return grpcUrl; + } + + public void setSecureGrpcUrl(String url) { + this.secureGrpcUrl = url; + } + + public String getSecureGrpcUrl() { + return secureGrpcUrl; + } } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/RequestTopicForPushRequest.java b/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/RequestTopicForPushRequest.java new file mode 100644 index 00000000000..580ee2de74c --- /dev/null +++ b/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/RequestTopicForPushRequest.java @@ -0,0 +1,178 @@ +package com.linkedin.venice.controllerapi; + +import com.linkedin.venice.meta.Version.PushType; +import java.security.cert.X509Certificate; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import java.util.Set; + + +public class RequestTopicForPushRequest { + private final String clusterName; + private final String storeName; + private final PushType pushType; + private final String pushJobId; + + private boolean sendStartOfPush = false; + private boolean sorted = false; // an inefficient but safe default + private boolean isWriteComputeEnabled = false; + private boolean separateRealTimeTopicEnabled = false; + private long rewindTimeInSecondsOverride = -1L; + private boolean deferVersionSwap = false; + private String targetedRegions = null; + private int repushSourceVersion = -1; + private Set partitioners = Collections.emptySet(); + private String compressionDictionary = null; + private X509Certificate certificateInRequest = null; + private String sourceGridFabric = null; + private String emergencySourceRegion = null; + + public RequestTopicForPushRequest(String clusterName, String storeName, PushType pushType, String pushJobId) { + if (clusterName == null || clusterName.isEmpty()) { + throw new IllegalArgumentException("clusterName is required"); + } + if (storeName == null || storeName.isEmpty()) { + throw new IllegalArgumentException("storeName is required"); + } + if (pushType == null) { + throw new IllegalArgumentException("pushType is required"); + } + + if (pushJobId == null || pushJobId.isEmpty()) { + throw new IllegalArgumentException("pushJobId is required"); + } + + this.clusterName = clusterName; + this.storeName = storeName; + this.pushType = pushType; + this.pushJobId = pushJobId; + } + + public String getClusterName() { + return clusterName; + } + + public String getStoreName() { + return storeName; + } + + public PushType getPushType() { + return pushType; + } + + public String getPushJobId() { + return pushJobId; + } + + public boolean isSendStartOfPush() { + return sendStartOfPush; + } + + public boolean isSorted() { + return sorted; + } + + public boolean isWriteComputeEnabled() { + return isWriteComputeEnabled; + } + + public String getSourceGridFabric() { + return sourceGridFabric; + } + + public long getRewindTimeInSecondsOverride() { + return rewindTimeInSecondsOverride; + } + + public boolean isDeferVersionSwap() { + return deferVersionSwap; + } + + public String getTargetedRegions() { + return targetedRegions; + } + + public int getRepushSourceVersion() { + return repushSourceVersion; + } + + public Set getPartitioners() { + return partitioners; + } + + public String getCompressionDictionary() { + return compressionDictionary; + } + + public X509Certificate getCertificateInRequest() { + return certificateInRequest; + } + + public String getEmergencySourceRegion() { + return emergencySourceRegion; + } + + public void setSendStartOfPush(boolean sendStartOfPush) { + this.sendStartOfPush = sendStartOfPush; + } + + public void setSorted(boolean sorted) { + this.sorted = sorted; + } + + public void setWriteComputeEnabled(boolean writeComputeEnabled) { + isWriteComputeEnabled = writeComputeEnabled; + } + + public void setSourceGridFabric(String sourceGridFabric) { + this.sourceGridFabric = sourceGridFabric; + } + + public void setRewindTimeInSecondsOverride(long rewindTimeInSecondsOverride) { + this.rewindTimeInSecondsOverride = rewindTimeInSecondsOverride; + } + + public void setDeferVersionSwap(boolean deferVersionSwap) { + this.deferVersionSwap = deferVersionSwap; + } + + public void setTargetedRegions(String targetedRegions) { + this.targetedRegions = targetedRegions; + } + + public void setRepushSourceVersion(int repushSourceVersion) { + this.repushSourceVersion = repushSourceVersion; + } + + public void setPartitioners(String commaSeparatedPartitioners) { + if (commaSeparatedPartitioners == null || commaSeparatedPartitioners.isEmpty()) { + return; + } + setPartitioners(new HashSet<>(Arrays.asList(commaSeparatedPartitioners.split(",")))); + } + + public void setPartitioners(Set partitioners) { + this.partitioners = partitioners; + } + + public void setCompressionDictionary(String compressionDictionary) { + this.compressionDictionary = compressionDictionary; + } + + public void setCertificateInRequest(X509Certificate certificateInRequest) { + this.certificateInRequest = certificateInRequest; + } + + public void setEmergencySourceRegion(String emergencySourceRegion) { + this.emergencySourceRegion = emergencySourceRegion; + } + + public boolean isSeparateRealTimeTopicEnabled() { + return separateRealTimeTopicEnabled; + } + + public void setSeparateRealTimeTopicEnabled(boolean separateRealTimeTopicEnabled) { + this.separateRealTimeTopicEnabled = separateRealTimeTopicEnabled; + } +} diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/transport/GrpcRequestResponseConverter.java b/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/transport/GrpcRequestResponseConverter.java new file mode 100644 index 00000000000..aeaedd69126 --- /dev/null +++ b/internal/venice-common/src/main/java/com/linkedin/venice/controllerapi/transport/GrpcRequestResponseConverter.java @@ -0,0 +1,127 @@ +package com.linkedin.venice.controllerapi.transport; + +import com.google.protobuf.Any; +import com.linkedin.venice.client.exceptions.VeniceClientException; +import com.linkedin.venice.controllerapi.ControllerResponse; +import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo; +import com.linkedin.venice.protocols.controller.ControllerGrpcErrorType; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcErrorInfo; +import io.grpc.Status.Code; +import io.grpc.StatusRuntimeException; +import io.grpc.protobuf.StatusProto; +import io.grpc.stub.StreamObserver; + + +public class GrpcRequestResponseConverter { + public static ClusterStoreGrpcInfo getClusterStoreGrpcInfo(ControllerResponse response) { + ClusterStoreGrpcInfo.Builder builder = ClusterStoreGrpcInfo.newBuilder(); + if (response.getCluster() != null) { + builder.setClusterName(response.getCluster()); + } + if (response.getName() != null) { + builder.setStoreName(response.getName()); + } + return builder.build(); + } + + /** + * Sends an error response to the client using the provided error details. + * + *

This method constructs a detailed gRPC error response, including the error type, status code, + * message, and optional cluster or store-specific information. The constructed error is sent + * to the client via the provided {@link StreamObserver}.

+ * + * @param code The gRPC status code representing the error (e.g., {@link io.grpc.Status.Code}). + * @param errorType The specific controller error type represented by {@link ControllerGrpcErrorType}. + * @param e The exception containing the error message. + * @param clusterName The name of the cluster associated with the error (can be null). + * @param storeName The name of the store associated with the error (can be null). + * @param responseObserver The {@link StreamObserver} to send the error response back to the client. + * + *

Example usage:

+ * 
+   * {@code
+   * sendErrorResponse(
+   *     Status.Code.INTERNAL,
+   *     ControllerGrpcErrorType.UNKNOWN_ERROR,
+   *     new RuntimeException("Something went wrong"),
+   *     "test-cluster",
+   *     "test-store",
+   *     responseObserver);
+   * }
+   *
+ * + *

The error response includes the following:

+ *
    + *
  • gRPC status code (e.g., INTERNAL, FAILED_PRECONDITION).
    • + *
  • Error type (e.g., BAD_REQUEST, CONCURRENT_BATCH_PUSH).
    • + *
  • Error message extracted from the provided exception.
    • + *
  • Optional cluster name and store name if provided.
    • + *
+ */ + public static void sendErrorResponse( + Code code, + ControllerGrpcErrorType errorType, + Exception e, + String clusterName, + String storeName, + StreamObserver responseObserver) { + VeniceControllerGrpcErrorInfo.Builder errorInfoBuilder = + VeniceControllerGrpcErrorInfo.newBuilder().setStatusCode(code.value()).setErrorType(errorType); + if (e.getMessage() != null) { + errorInfoBuilder.setErrorMessage(e.getMessage()); + } + if (clusterName != null) { + errorInfoBuilder.setClusterName(clusterName); + } + if (storeName != null) { + errorInfoBuilder.setStoreName(storeName); + } + // Wrap the error info into a com.google.rpc.Status message + com.google.rpc.Status status = + com.google.rpc.Status.newBuilder().setCode(code.value()).addDetails(Any.pack(errorInfoBuilder.build())).build(); + + // Send the error response + responseObserver.onError(StatusProto.toStatusRuntimeException(status)); + } + + /** + * Parses a {@link StatusRuntimeException} to extract a {@link VeniceControllerGrpcErrorInfo} object. + * + *

This method processes the gRPC error details embedded within a {@link StatusRuntimeException}. + * If the error details contain a {@link VeniceControllerGrpcErrorInfo}, it unpacks and returns it. + * If no valid details are found or the unpacking fails, a {@link VeniceClientException} is thrown.

+ * + * @param e The {@link StatusRuntimeException} containing the gRPC error details. + * @return A {@link VeniceControllerGrpcErrorInfo} object extracted from the error details. + * @throws VeniceClientException If the error details cannot be unpacked or no valid information is found. + * + *

Example usage:

+ * 
+   * {@code
+   * try {
+   *     // Call a gRPC method that might throw StatusRuntimeException
+   * } catch (StatusRuntimeException e) {
+   *     VeniceControllerGrpcErrorInfo errorInfo = parseControllerGrpcError(e);
+   *     System.out.println("Error Type: " + errorInfo.getErrorType());
+   *     System.out.println("Error Message: " + errorInfo.getErrorMessage());
+   * }
+   * }
+   *
+ */ + public static VeniceControllerGrpcErrorInfo parseControllerGrpcError(StatusRuntimeException e) { + com.google.rpc.Status status = StatusProto.fromThrowable(e); + if (status != null) { + for (com.google.protobuf.Any detail: status.getDetailsList()) { + if (detail.is(VeniceControllerGrpcErrorInfo.class)) { + try { + return detail.unpack(VeniceControllerGrpcErrorInfo.class); + } catch (Exception unpackException) { + throw new VeniceClientException("Failed to unpack error details: " + unpackException.getMessage()); + } + } + } + } + throw new VeniceClientException("An unknown gRPC error occurred. Error code: " + Code.UNKNOWN.name()); + } +} diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/grpc/GrpcUtils.java b/internal/venice-common/src/main/java/com/linkedin/venice/grpc/GrpcUtils.java index f2fcb1c0e7a..818d7056713 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/grpc/GrpcUtils.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/grpc/GrpcUtils.java @@ -1,13 +1,17 @@ package com.linkedin.venice.grpc; import com.linkedin.venice.acl.handler.AccessResult; +import com.linkedin.venice.client.exceptions.VeniceClientException; import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.security.SSLConfig; import com.linkedin.venice.security.SSLFactory; import com.linkedin.venice.utils.SslUtils; +import io.grpc.ChannelCredentials; import io.grpc.Grpc; +import io.grpc.InsecureChannelCredentials; import io.grpc.ServerCall; import io.grpc.Status; +import io.grpc.TlsChannelCredentials; import java.io.IOException; import java.io.InputStream; import java.nio.file.Files; @@ -97,4 +101,22 @@ private static KeyStore loadStore(String path, char[] password, String type) } return keyStore; } + + public static ChannelCredentials buildChannelCredentials(SSLFactory sslFactory) { + // TODO: Evaluate if this needs to fail instead since it depends on plain text support on server + if (sslFactory == null) { + return InsecureChannelCredentials.create(); + } + + try { + TlsChannelCredentials.Builder tlsBuilder = TlsChannelCredentials.newBuilder() + .keyManager(GrpcUtils.getKeyManagers(sslFactory)) + .trustManager(GrpcUtils.getTrustManagers(sslFactory)); + return tlsBuilder.build(); + } catch (Exception e) { + throw new VeniceClientException( + "Failed to initialize SSL channel credentials for Venice gRPC Transport Client", + e); + } + } } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServer.java b/internal/venice-common/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServer.java similarity index 74% rename from services/venice-server/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServer.java rename to internal/venice-common/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServer.java index 3e6fc3546a1..0122c396e23 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServer.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServer.java @@ -26,18 +26,16 @@ public class VeniceGrpcServer { private final SSLFactory sslFactory; // protected for testing purposes protected ServerCredentials credentials; + boolean isSecure; private final Executor executor; private final VeniceGrpcServerConfig config; public VeniceGrpcServer(VeniceGrpcServerConfig config) { - port = config.getPort(); - sslFactory = config.getSslFactory(); - executor = config.getExecutor(); - + this.port = config.getPort(); + this.sslFactory = config.getSslFactory(); + this.executor = config.getExecutor(); this.config = config; - initServerCredentials(); - server = Grpc.newServerBuilderForPort(config.getPort(), credentials) .executor(executor) // TODO: experiment with different executors for best performance .addService(ServerInterceptors.intercept(config.getService(), config.getInterceptors())) @@ -47,14 +45,16 @@ public VeniceGrpcServer(VeniceGrpcServerConfig config) { private void initServerCredentials() { if (sslFactory == null && config.getCredentials() == null) { - LOGGER.info("Creating gRPC server with insecure credentials"); + LOGGER.info("Creating gRPC server with insecure credentials on port: {}", port); credentials = InsecureServerCredentials.create(); + isSecure = false; return; } if (config.getCredentials() != null) { - LOGGER.info("Creating gRPC server with custom credentials"); + LOGGER.debug("Creating gRPC server with custom credentials"); credentials = config.getCredentials(); + isSecure = !(credentials instanceof InsecureServerCredentials); return; } @@ -64,6 +64,7 @@ private void initServerCredentials() { .trustManager(GrpcUtils.getTrustManagers(sslFactory)) .clientAuth(TlsServerCredentials.ClientAuth.REQUIRE) .build(); + isSecure = !(credentials instanceof InsecureServerCredentials); } catch (UnrecoverableKeyException | KeyStoreException | CertificateException | IOException | NoSuchAlgorithmException e) { LOGGER.error("Failed to initialize secure server credentials for gRPC Server"); @@ -74,9 +75,14 @@ private void initServerCredentials() { public void start() throws VeniceException { try { server.start(); + LOGGER.info( + "Started gRPC server for service: {} on port: {} isSecure: {}", + config.getService().getClass().getSimpleName(), + port, + isSecure()); } catch (IOException exception) { LOGGER.error( - "Failed to start gRPC Server for service {} on port {}", + "Failed to start gRPC server for service: {} on port: {}", config.getService().getClass().getSimpleName(), port, exception); @@ -84,17 +90,30 @@ public void start() throws VeniceException { } } - public boolean isShutdown() { - return server.isShutdown(); + public boolean isRunning() { + return !server.isShutdown(); } public boolean isTerminated() { return server.isTerminated(); } + private boolean isSecure() { + return isSecure; + } + public void stop() { + LOGGER.info( + "Shutting down gRPC server for service: {} on port: {} isSecure: {}", + config.getService().getClass().getSimpleName(), + port, + isSecure()); if (server != null && !server.isShutdown()) { server.shutdown(); } } + + public Server getServer() { + return server; + } } diff --git a/services/venice-server/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServerConfig.java b/internal/venice-common/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServerConfig.java similarity index 89% rename from services/venice-server/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServerConfig.java rename to internal/venice-common/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServerConfig.java index c0a2f5b9841..96b271f1f27 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServerConfig.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/grpc/VeniceGrpcServerConfig.java @@ -112,18 +112,18 @@ public VeniceGrpcServerConfig build() { private void verifyAndAddDefaults() { if (port == null) { - throw new IllegalArgumentException("Port must be set"); + throw new IllegalArgumentException("Port value is required to create the gRPC server but was not provided."); } if (service == null) { - throw new IllegalArgumentException("Service must be set"); + throw new IllegalArgumentException("A non-null gRPC service instance is required to create the server."); + } + if (numThreads <= 0 && executor == null) { + throw new IllegalArgumentException( + "gRPC server creation requires a valid number of threads (numThreads > 0) or a non-null executor."); } if (interceptors == null) { interceptors = Collections.emptyList(); } - if (numThreads <= 0 && executor == null) { - throw new IllegalArgumentException("Either numThreads or executor must be set"); - } - if (executor == null) { executor = Executors.newFixedThreadPool(numThreads); } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/meta/Instance.java b/internal/venice-common/src/main/java/com/linkedin/venice/meta/Instance.java index 7b0aa5c7a3f..29a848d5029 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/meta/Instance.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/meta/Instance.java @@ -34,23 +34,38 @@ public class Instance { private final String url; private final String sUrl; + private final int grpcPort; + private final int grpcSslPort; + private final String grpcUrl; + private final String grpcSslUrl; + // TODO: generate nodeId from host and port, should be "host_port", or generate host and port from id. public Instance(String nodeId, String host, int port) { - this(nodeId, host, port, port); + this(nodeId, host, port, port, -1, -1); + } + + public Instance(String nodeId, String host, int port, int grpcPort, int grpcSslPort) { + this(nodeId, host, port, port, grpcPort, grpcSslPort); } public Instance( @JsonProperty("nodeId") String nodeId, @JsonProperty("host") String host, @JsonProperty("port") int port, - @JsonProperty("sslPort") int sslPort) { + @JsonProperty("sslPort") int sslPort, + @JsonProperty("grpcPort") int grpcPort, + @JsonProperty("grpcSslPort") int grpcSslPort) { this.nodeId = nodeId; this.host = host; - validatePort("port", port); - this.port = port; + this.port = validatePort("port", port); this.sslPort = sslPort; this.url = "http://" + host + ":" + port + "/"; this.sUrl = "https://" + host + ":" + sslPort + "/"; + + this.grpcPort = grpcPort; + this.grpcSslPort = grpcSslPort; + this.grpcUrl = host + ":" + grpcPort; + this.grpcSslUrl = host + ":" + grpcSslPort; } public static Instance fromHostAndPort(String hostName, int port) { @@ -91,6 +106,22 @@ public int getSslPort() { return sslPort; } + public int getGrpcPort() { + return grpcPort; + } + + public int getGrpcSslPort() { + return grpcSslPort; + } + + public String getGrpcUrl() { + return grpcUrl; + } + + public String getGrpcSslUrl() { + return grpcSslUrl; + } + /*** * Convenience method for getting a host and port based url. * Wraps IPv6 host strings in square brackets @@ -111,10 +142,11 @@ public String getUrl() { return getUrl(false); } - private void validatePort(String name, int port) { + private int validatePort(String name, int port) { if (port < 0 || port > 65535) { throw new IllegalArgumentException("Invalid " + name + ": " + port); } + return port; } // Autogen except for .toLowerCase() diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/meta/NameRepository.java b/internal/venice-common/src/main/java/com/linkedin/venice/meta/NameRepository.java new file mode 100644 index 00000000000..2bc17d6fe13 --- /dev/null +++ b/internal/venice-common/src/main/java/com/linkedin/venice/meta/NameRepository.java @@ -0,0 +1,39 @@ +package com.linkedin.venice.meta; + +import com.github.benmanes.caffeine.cache.Caffeine; +import com.github.benmanes.caffeine.cache.LoadingCache; + + +/** + * A repository of {@link StoreName} and {@link StoreVersionName}, which are intended to be used as shared instances. + */ +public class NameRepository { + private static final int DEFAULT_MAXIMUM_ENTRY_COUNT = 2000; + private final LoadingCache storeNameCache; + private final LoadingCache storeVersionNameCache; + + public NameRepository() { + this(DEFAULT_MAXIMUM_ENTRY_COUNT); + } + + public NameRepository(int maxEntryCount) { + this.storeNameCache = Caffeine.newBuilder().maximumSize(maxEntryCount).build(StoreName::new); + this.storeVersionNameCache = Caffeine.newBuilder().maximumSize(maxEntryCount).build(storeVersionNameString -> { + String storeNameString = Version.parseStoreFromKafkaTopicName(storeVersionNameString); + StoreName storeName = getStoreName(storeNameString); + return new StoreVersionName(storeVersionNameString, storeName); + }); + } + + public StoreName getStoreName(String storeName) { + return this.storeNameCache.get(storeName); + } + + public StoreVersionName getStoreVersionName(String storeVersionName) { + return this.storeVersionNameCache.get(storeVersionName); + } + + public StoreVersionName getStoreVersionName(String storeName, int versionNumber) { + return getStoreVersionName(Version.composeKafkaTopic(storeName, versionNumber)); + } +} diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/meta/Store.java b/internal/venice-common/src/main/java/com/linkedin/venice/meta/Store.java index 03a4b7f9261..e0946e9271e 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/meta/Store.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/meta/Store.java @@ -6,8 +6,6 @@ import java.util.List; import java.util.Map; import java.util.concurrent.TimeUnit; -import java.util.regex.Matcher; -import java.util.regex.Pattern; import javax.annotation.Nonnull; import javax.annotation.Nullable; @@ -43,19 +41,6 @@ public interface Store { int DEFAULT_BATCH_GET_LIMIT = 150; - /** - * Store name rules: - * 1. Only letters, numbers, underscore or dash - * 2. No double dashes - */ - - Pattern storeNamePattern = Pattern.compile("^[a-zA-Z0-9_-]+$"); - - static boolean isValidStoreName(String name) { - Matcher matcher = storeNamePattern.matcher(name); - return matcher.matches() && !name.contains("--"); - } - static boolean isSystemStore(String storeName) { return storeName.startsWith(SYSTEM_STORE_NAME_PREFIX); } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/meta/StoreName.java b/internal/venice-common/src/main/java/com/linkedin/venice/meta/StoreName.java new file mode 100644 index 00000000000..1d1f1aa3e98 --- /dev/null +++ b/internal/venice-common/src/main/java/com/linkedin/venice/meta/StoreName.java @@ -0,0 +1,76 @@ +package com.linkedin.venice.meta; + +import java.util.Objects; +import java.util.regex.Matcher; +import java.util.regex.Pattern; +import javax.annotation.Nonnull; + + +/** + * This class is a handle to refer to a store. + * + * It intentionally does not contain any operational state related to the referenced store. + * + * It is appropriate to use as a map key. Its {@link #equals(Object)} and {@link #hashCode()} delegate to the same + * functions on the String form of the store name. + * + * The purpose of using this handle class rather than a String is two-fold: + * + * - It is a stronger type than String, since a String can contain anything. + * - It can be more performant, since shared instances are allocated once and reused thus causing less garbage, and + * shared instances are also faster to use as map keys (the hash code gets cached, and equality checks can be resolved + * by identity, in the common case). + */ +public class StoreName { + private static final Pattern STORE_NAME_PATTERN = Pattern.compile("^[a-zA-Z0-9_-]+$"); + + private final String name; + + /** + * Not intended to be called directly! Instead, use: + * + * {@link NameRepository#getStoreName(String)} + */ + StoreName(String name) { + if (!isValidStoreName(name)) { + throw new IllegalArgumentException("Invalid store name!"); + } + this.name = Objects.requireNonNull(name); + } + + /** + * Store name rules: + * 1. Only letters, numbers, underscore or dash + * 2. No double dashes + */ + public static boolean isValidStoreName(String name) { + Matcher matcher = STORE_NAME_PATTERN.matcher(name); + return matcher.matches() && !name.contains("--"); + } + + @Nonnull + public String getName() { + return this.name; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } else if (o instanceof StoreName) { + StoreName that = (StoreName) o; + return this.name.equals(that.name); + } + return false; + } + + @Override + public int hashCode() { + return this.name.hashCode(); + } + + @Override + public String toString() { + return this.getClass().getSimpleName() + "(" + this.name + ")"; + } +} diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/meta/StoreVersionName.java b/internal/venice-common/src/main/java/com/linkedin/venice/meta/StoreVersionName.java new file mode 100644 index 00000000000..a188ba90658 --- /dev/null +++ b/internal/venice-common/src/main/java/com/linkedin/venice/meta/StoreVersionName.java @@ -0,0 +1,82 @@ +package com.linkedin.venice.meta; + +import java.util.Objects; +import javax.annotation.Nonnull; + + +/** + * This class is a handle to refer to a store-version. + * + * It intentionally does not contain any operational state related to the referenced store-version. + * + * It is appropriate to use as a map key. Its {@link #equals(Object)} and {@link #hashCode()} delegate to the same + * functions on the String form of the store-version name. + * + * The purpose of using this handle class rather than a String is two-fold: + * + * - It is a stronger type than String, since a String can contain anything. + * - It can be more performant, since shared instances are allocated once and reused thus causing less garbage, and + * shared instances are also faster to use as map keys (the hash code gets cached, and equality checks can be resolved + * by identity, in the common case). + */ +public class StoreVersionName { + private final String name; + private final StoreName storeName; + private final int versionNumber; + + /** + * Not intended to be called directly! Instead, use: + * + * {@link NameRepository#getStoreVersionName(String)} + * {@link NameRepository#getStoreVersionName(String, int)} + */ + StoreVersionName(String storeVersionName, StoreName storeName) { + if (!Version.isVersionTopic(storeVersionName)) { + throw new IllegalArgumentException( + "Only valid store-version names are allowed! Valid names take the form of: storeName + \"_v\" + versionNumber"); + } + this.name = Objects.requireNonNull(storeVersionName); + this.storeName = Objects.requireNonNull(storeName); + this.versionNumber = Version.parseVersionFromVersionTopicName(storeVersionName); + } + + @Nonnull + public String getName() { + return this.name; + } + + @Nonnull + public StoreName getStore() { + return this.storeName; + } + + @Nonnull + public String getStoreName() { + return this.storeName.getName(); + } + + public int getVersionNumber() { + return this.versionNumber; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } else if (o instanceof StoreVersionName) { + StoreVersionName that = (StoreVersionName) o; + return this.name.equals(that.name); + } + return false; + } + + @Override + public int hashCode() { + return this.name.hashCode(); + } + + @Override + public String toString() { + return this.getClass().getSimpleName() + "(" + this.name + ")"; + } +} diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/meta/Version.java b/internal/venice-common/src/main/java/com/linkedin/venice/meta/Version.java index c5092013141..7015dcce6e5 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/meta/Version.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/meta/Version.java @@ -10,10 +10,9 @@ import com.linkedin.venice.systemstore.schemas.StoreVersion; import com.linkedin.venice.views.VeniceView; import java.time.Duration; -import java.util.Arrays; +import java.util.HashMap; import java.util.List; import java.util.Map; -import java.util.Optional; /** @@ -43,11 +42,21 @@ public interface Version extends Comparable, DataModelBackedStructure VALUE_TO_TYPE_MAP = new HashMap<>(4); + private static final Map NAME_TO_TYPE_MAP = new HashMap<>(4); + + // Static initializer for map population + static { + for (PushType type: PushType.values()) { + VALUE_TO_TYPE_MAP.put(type.value, type); + NAME_TO_TYPE_MAP.put(type.name(), type); + } + } PushType(int value) { this.value = value; @@ -70,15 +79,41 @@ public boolean isStreamReprocessing() { } public boolean isBatchOrStreamReprocessing() { - return isBatch() || isStreamReprocessing(); + return this == BATCH || this == STREAM_REPROCESSING; } + /** + * Retrieve the PushType based on its integer value. + * + * @param value the integer value of the PushType + * @return the corresponding PushType + * @throws VeniceException if the value is invalid + */ public static PushType valueOf(int value) { - Optional pushType = Arrays.stream(values()).filter(p -> p.value == value).findFirst(); - if (!pushType.isPresent()) { + PushType pushType = VALUE_TO_TYPE_MAP.get(value); + if (pushType == null) { throw new VeniceException("Invalid push type with int value: " + value); } - return pushType.get(); + return pushType; + } + + /** + * Extracts the PushType from its string name. + * + * @param pushTypeString the string representation of the PushType + * @return the corresponding PushType + * @throws IllegalArgumentException if the string is invalid + */ + public static PushType extractPushType(String pushTypeString) { + PushType pushType = NAME_TO_TYPE_MAP.get(pushTypeString); + if (pushType == null) { + throw new IllegalArgumentException( + String.format( + "%s is an invalid push type. Valid push types are: %s", + pushTypeString, + String.join(", ", NAME_TO_TYPE_MAP.keySet()))); + } + return pushType; } } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/meta/ZKStore.java b/internal/venice-common/src/main/java/com/linkedin/venice/meta/ZKStore.java index 1a05c0a807b..44a39d589fb 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/meta/ZKStore.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/meta/ZKStore.java @@ -93,7 +93,7 @@ public ZKStore( HybridStoreConfig hybridStoreConfig, PartitionerConfig partitionerConfig, int replicationFactor) { - if (!Store.isValidStoreName(name)) { + if (!StoreName.isValidStoreName(name)) { throw new VeniceException("Invalid store name: " + name); } @@ -151,7 +151,7 @@ public List getForRead() { } public ZKStore(StoreProperties storeProperties) { - if (!Store.isValidStoreName(storeProperties.name.toString())) { + if (!StoreName.isValidStoreName(storeProperties.name.toString())) { throw new VeniceException("Invalid store name: " + storeProperties.name.toString()); } this.storeProperties = storeProperties; diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/PubSubConstants.java b/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/PubSubConstants.java index 011e468acb7..5423b3124b2 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/PubSubConstants.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/PubSubConstants.java @@ -2,6 +2,7 @@ import com.linkedin.venice.pubsub.api.exceptions.PubSubClientRetriableException; import com.linkedin.venice.pubsub.api.exceptions.PubSubOpTimeoutException; +import com.linkedin.venice.pubsub.api.exceptions.PubSubTopicDoesNotExistException; import com.linkedin.venice.utils.Time; import java.time.Duration; import java.util.Arrays; @@ -61,10 +62,18 @@ public class PubSubConstants { public static final List> CREATE_TOPIC_RETRIABLE_EXCEPTIONS = Collections.unmodifiableList(Arrays.asList(PubSubOpTimeoutException.class, PubSubClientRetriableException.class)); + public static final List> TOPIC_METADATA_OP_RETRIABLE_EXCEPTIONS = + Collections.unmodifiableList( + Arrays.asList( + PubSubTopicDoesNotExistException.class, + PubSubOpTimeoutException.class, + PubSubClientRetriableException.class)); + /** * Default value of sleep interval for polling topic deletion status from ZK. */ public static final int PUBSUB_TOPIC_DELETION_STATUS_POLL_INTERVAL_MS_DEFAULT_VALUE = 2 * Time.MS_PER_SECOND; + public static final long UNKNOWN_LATEST_OFFSET = -1L; private static final Duration PUBSUB_OFFSET_API_TIMEOUT_DURATION_DEFAULT_VALUE_DEFAULT = Duration.ofMinutes(1); private static Duration PUBSUB_OFFSET_API_TIMEOUT_DURATION_DEFAULT_VALUE = diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/PubSubTopicConfiguration.java b/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/PubSubTopicConfiguration.java index 02152127e3a..0a98549676b 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/PubSubTopicConfiguration.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/PubSubTopicConfiguration.java @@ -6,7 +6,7 @@ /** * Represents a {@link com.linkedin.venice.pubsub.api.PubSubTopic} configuration. */ -public class PubSubTopicConfiguration { +public class PubSubTopicConfiguration implements Cloneable { Optional retentionInMs; boolean isLogCompacted; Long minLogCompactionLagMs; @@ -103,4 +103,9 @@ public String toString() { minLogCompactionLagMs, maxLogCompactionLagMs.isPresent() ? maxLogCompactionLagMs.get() : " not set"); } + + @Override + public PubSubTopicConfiguration clone() throws CloneNotSupportedException { + return (PubSubTopicConfiguration) super.clone(); + } } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/manager/TopicManager.java b/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/manager/TopicManager.java index c91d3c452ef..13af9da7f2a 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/manager/TopicManager.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/manager/TopicManager.java @@ -6,6 +6,7 @@ import static com.linkedin.venice.pubsub.PubSubConstants.PUBSUB_FAST_OPERATION_TIMEOUT_MS; import static com.linkedin.venice.pubsub.PubSubConstants.PUBSUB_TOPIC_DELETE_RETRY_TIMES; import static com.linkedin.venice.pubsub.PubSubConstants.PUBSUB_TOPIC_UNKNOWN_RETENTION; +import static com.linkedin.venice.pubsub.PubSubConstants.TOPIC_METADATA_OP_RETRIABLE_EXCEPTIONS; import static com.linkedin.venice.pubsub.manager.TopicManagerStats.SENSOR_TYPE.CONTAINS_TOPIC_WITH_RETRY; import static com.linkedin.venice.pubsub.manager.TopicManagerStats.SENSOR_TYPE.CREATE_TOPIC; import static com.linkedin.venice.pubsub.manager.TopicManagerStats.SENSOR_TYPE.DELETE_TOPIC; @@ -272,6 +273,35 @@ public boolean updateTopicRetention( return false; } + public boolean updateTopicRetentionWithRetries(PubSubTopic topicName, long expectedRetentionInMs) { + PubSubTopicConfiguration topicConfiguration; + try { + topicConfiguration = getCachedTopicConfig(topicName).clone(); + } catch (Exception e) { + logger.error("Failed to get topic config for topic: {}", topicName, e); + throw new VeniceException( + "Failed to update topic retention for topic: " + topicName + " with retention: " + expectedRetentionInMs + + " in cluster: " + this.pubSubClusterAddress, + e); + } + if (topicConfiguration.retentionInMs().isPresent() + && topicConfiguration.retentionInMs().get() == expectedRetentionInMs) { + // Retention time has already been updated for this topic before + return false; + } + + topicConfiguration.setRetentionInMs(Optional.of(expectedRetentionInMs)); + RetryUtils.executeWithMaxAttemptAndExponentialBackoff( + () -> setTopicConfig(topicName, topicConfiguration), + 5, + Duration.ofMillis(200), + Duration.ofSeconds(1), + Duration.ofMinutes(2), + TOPIC_METADATA_OP_RETRIABLE_EXCEPTIONS); + topicConfigCache.put(topicName, topicConfiguration); + return true; + } + public void updateTopicCompactionPolicy(PubSubTopic topic, boolean expectedLogCompacted) { updateTopicCompactionPolicy(topic, expectedLogCompacted, -1, Optional.empty()); } @@ -436,7 +466,7 @@ private void setTopicConfig(PubSubTopic pubSubTopic, PubSubTopicConfiguration pu pubSubAdminAdapter.setTopicConfig(pubSubTopic, pubSubTopicConfiguration); stats.recordLatency(SET_TOPIC_CONFIG, startTime); } catch (Exception e) { - logger.debug("Failed to set topic config for topic: {}", pubSubTopic, e); + logger.info("Failed to set topic config for topic: {}", pubSubTopic, e); stats.recordPubSubAdminOpFailure(); throw e; } diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/manager/TopicMetadataFetcher.java b/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/manager/TopicMetadataFetcher.java index 4a4e7ac85e4..e34602e4733 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/manager/TopicMetadataFetcher.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/pubsub/manager/TopicMetadataFetcher.java @@ -392,7 +392,7 @@ long getLatestOffsetCachedNonBlocking(PubSubTopicPartition pubSubTopicPartition) if (cachedValue == null) { cachedValue = latestOffsetCache.get(pubSubTopicPartition); if (cachedValue == null) { - return -1; + return PubSubConstants.UNKNOWN_LATEST_OFFSET; } } return cachedValue.getValue(); diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/throttle/EventThrottler.java b/internal/venice-common/src/main/java/com/linkedin/venice/throttle/EventThrottler.java index 93dd8365db0..24e148fc16a 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/throttle/EventThrottler.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/throttle/EventThrottler.java @@ -54,6 +54,10 @@ public class EventThrottler implements VeniceRateLimiter { // Used only to compare if the new quota requests are different from the existing quota. private long quota = -1; + public EventThrottler() { + this(1, DEFAULT_CHECK_INTERVAL_MS, null, false, BLOCK_STRATEGY); + } + /** * @param maxRatePerSecond Maximum rate that this throttler should allow (-1 is unlimited) */ diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/utils/Utils.java b/internal/venice-common/src/main/java/com/linkedin/venice/utils/Utils.java index 72b8fe5e119..11e10696296 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/utils/Utils.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/utils/Utils.java @@ -73,7 +73,6 @@ import org.apache.http.HttpStatus; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; -import org.apache.logging.log4j.util.Strings; /** @@ -309,6 +308,12 @@ public static long parseLongFromString(String value, String fieldName) { * any string that are not equal to 'true', We validate the string by our own. */ public static boolean parseBooleanFromString(String value, String fieldName) { + if (value == null) { + throw new VeniceHttpException( + HttpStatus.SC_BAD_REQUEST, + fieldName + " must be a boolean, but value is null", + ErrorType.BAD_REQUEST); + } if (value.equalsIgnoreCase("true") || value.equalsIgnoreCase("false")) { return Boolean.parseBoolean(value); } else { @@ -545,6 +550,18 @@ public static String composeRealTimeTopic(String storeName) { return storeName + Version.REAL_TIME_TOPIC_SUFFIX; } + public static String getRealTimeTopicNameFromStoreConfig(Store store) { + HybridStoreConfig hybridStoreConfig = store.getHybridStoreConfig(); + String storeName = store.getName(); + + if (hybridStoreConfig != null) { + String realTimeTopicName = hybridStoreConfig.getRealTimeTopicName(); + return getRealTimeTopicNameIfEmpty(realTimeTopicName, storeName); + } else { + return composeRealTimeTopic(storeName); + } + } + /** * It follows the following order to search for real time topic name, * i) current store-version config, ii) store config, iii) other store-version configs, iv) default name @@ -590,7 +607,7 @@ static String getRealTimeTopicName( versions.stream().filter(version -> version.getNumber() == currentVersionNumber).findFirst(); if (currentVersion.isPresent() && currentVersion.get().isHybrid()) { String realTimeTopicName = currentVersion.get().getHybridStoreConfig().getRealTimeTopicName(); - if (Strings.isNotBlank(realTimeTopicName)) { + if (StringUtils.isNotBlank(realTimeTopicName)) { return realTimeTopicName; } } @@ -606,7 +623,7 @@ static String getRealTimeTopicName( try { if (version.isHybrid()) { String realTimeTopicName = version.getHybridStoreConfig().getRealTimeTopicName(); - if (Strings.isNotBlank(realTimeTopicName)) { + if (StringUtils.isNotBlank(realTimeTopicName)) { realTimeTopicNames.add(realTimeTopicName); } } @@ -629,7 +646,31 @@ static String getRealTimeTopicName( } private static String getRealTimeTopicNameIfEmpty(String realTimeTopicName, String storeName) { - return Strings.isBlank(realTimeTopicName) ? composeRealTimeTopic(storeName) : realTimeTopicName; + return StringUtils.isBlank(realTimeTopicName) ? composeRealTimeTopic(storeName) : realTimeTopicName; + } + + public static String createNewRealTimeTopicName(String oldRealTimeTopicName) { + if (oldRealTimeTopicName == null || !oldRealTimeTopicName.endsWith(Version.REAL_TIME_TOPIC_SUFFIX)) { + throw new IllegalArgumentException("Invalid old name format"); + } + + // Extract the base name and current version + int suffixLength = Version.REAL_TIME_TOPIC_SUFFIX.length(); + String base = oldRealTimeTopicName.substring(0, oldRealTimeTopicName.length() - suffixLength); + + // Locate the last version separator "_v" in the base + int versionSeparatorIndex = base.lastIndexOf("_v"); + if (versionSeparatorIndex > -1 && versionSeparatorIndex < base.length() - 2) { + // Extract and increment the version + String versionStr = base.substring(versionSeparatorIndex + 2); + int version = Integer.parseInt(versionStr) + 1; + base = base.substring(0, versionSeparatorIndex) + "_v" + version; + } else { + // Start with version 2 if no valid version is present + base = base + "_v2"; + } + + return base + Version.REAL_TIME_TOPIC_SUFFIX; } private static class TimeUnitInfo { diff --git a/internal/venice-common/src/main/java/com/linkedin/venice/writer/VeniceWriter.java b/internal/venice-common/src/main/java/com/linkedin/venice/writer/VeniceWriter.java index 5d8cd79cb28..c8bce097c4a 100644 --- a/internal/venice-common/src/main/java/com/linkedin/venice/writer/VeniceWriter.java +++ b/internal/venice-common/src/main/java/com/linkedin/venice/writer/VeniceWriter.java @@ -1895,6 +1895,12 @@ public CompletableFuture sendHeartbeat( boolean addLeaderCompleteState, LeaderCompleteState leaderCompleteState, long originTimeStampMs) { + if (isClosed) { + CompletableFuture future = new CompletableFuture<>(); + future.completedFuture(null); + logger.warn("VeniceWriter already closed for topic partition " + topicPartition); + return future; + } KafkaMessageEnvelope kafkaMessageEnvelope = getHeartbeatKME(originTimeStampMs, leaderMetadataWrapper, heartBeatMessage, writerId); return producerAdapter.sendMessage( diff --git a/internal/venice-common/src/main/proto/VeniceControllerGrpcService.proto b/internal/venice-common/src/main/proto/VeniceControllerGrpcService.proto new file mode 100644 index 00000000000..eedba3e0e32 --- /dev/null +++ b/internal/venice-common/src/main/proto/VeniceControllerGrpcService.proto @@ -0,0 +1,85 @@ +syntax = 'proto3'; +package com.linkedin.venice.protocols.controller; + +import "google/rpc/status.proto"; +import "google/rpc/error_details.proto"; + +option java_multiple_files = true; + + +service VeniceControllerGrpcService { + // ClusterDiscovery + rpc discoverClusterForStore(DiscoverClusterGrpcRequest) returns (DiscoverClusterGrpcResponse) {} + + // ControllerRoutes + rpc getLeaderController(LeaderControllerGrpcRequest) returns (LeaderControllerGrpcResponse); + + // CreateStore + rpc createStore(CreateStoreGrpcRequest) returns (CreateStoreGrpcResponse) {} +} + +message DiscoverClusterGrpcRequest { + string storeName = 1; +} + +message DiscoverClusterGrpcResponse { + string clusterName = 1; + string storeName = 2; + string d2Service = 3; + string serverD2Service = 4; + optional string zkAddress = 5; + optional string pubSubBootstrapServers = 6; +} + +message ClusterStoreGrpcInfo { + string clusterName = 1; + string storeName = 2; +} + +message CreateStoreGrpcRequest { + ClusterStoreGrpcInfo clusterStoreInfo = 1; + string keySchema = 2; + string valueSchema = 3; + optional string owner = 4; + optional bool isSystemStore = 5; + optional string accessPermission = 6; +} + +message CreateStoreGrpcResponse { + ClusterStoreGrpcInfo clusterStoreInfo = 1; + string owner = 2; +} + +enum ControllerGrpcErrorType { + UNKNOWN = 0; + INCORRECT_CONTROLLER = 1; + INVALID_SCHEMA = 2; + INVALID_CONFIG = 3; + STORE_NOT_FOUND = 4; + SCHEMA_NOT_FOUND = 5; + CONNECTION_ERROR = 6; + GENERAL_ERROR = 7; + BAD_REQUEST = 8; + CONCURRENT_BATCH_PUSH = 9; + RESOURCE_STILL_EXISTS = 10; +} + +message VeniceControllerGrpcErrorInfo { + uint32 statusCode = 1; + string errorMessage = 2; + optional ControllerGrpcErrorType errorType = 3; + optional string clusterName = 4; + optional string storeName = 5; +} + +message LeaderControllerGrpcRequest { + string clusterName = 1; // The cluster name +} + +message LeaderControllerGrpcResponse { + string clusterName = 1; // The cluster name + string httpUrl = 2; // Leader controller URL + string httpsUrl = 3; // SSL-enabled leader controller URL + string grpcUrl = 4; // gRPC URL for leader controller + string secureGrpcUrl = 5; // Secure gRPC URL for leader controller +} diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/acl/NoOpDynamicAccessControllerTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/acl/NoOpDynamicAccessControllerTest.java new file mode 100644 index 00000000000..8f737ff0dc5 --- /dev/null +++ b/internal/venice-common/src/test/java/com/linkedin/venice/acl/NoOpDynamicAccessControllerTest.java @@ -0,0 +1,67 @@ +package com.linkedin.venice.acl; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertTrue; + +import java.security.cert.X509Certificate; +import java.util.Collections; +import java.util.Set; +import javax.security.auth.x500.X500Principal; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class NoOpDynamicAccessControllerTest { + private NoOpDynamicAccessController accessController; + + @BeforeMethod + public void setUp() { + accessController = NoOpDynamicAccessController.INSTANCE; + } + + @Test + public void testAlwaysTrueMethods() throws AclException { + // Test all methods that always return true + assertTrue(accessController.hasAccess(null, "resource", "GET"), "Expected hasAccess to return true."); + assertTrue(accessController.hasAccessToTopic(null, "topic", "READ"), "Expected hasAccessToTopic to return true."); + assertTrue( + accessController.hasAccessToAdminOperation(null, "operation"), + "Expected hasAccessToAdminOperation to return true."); + assertTrue(accessController.isAllowlistUsers(null, "resource", "GET"), "Expected isAllowlistUsers to return true."); + assertTrue(accessController.hasAcl("resource"), "Expected hasAcl to return true."); + } + + @Test + public void testInitReturnsSameInstance() { + DynamicAccessController initializedController = accessController.init(Collections.emptyList()); + assertSame(initializedController, accessController, "Expected the same instance after init."); + } + + @Test + public void testGetPrincipalId() { + String expectedId = "CN=Test User,OU=Engineering,O=LinkedIn,C=US"; + X509Certificate mockCertificate = mock(X509Certificate.class); + when(mockCertificate.getSubjectX500Principal()).thenReturn(new X500Principal(expectedId)); + + assertEquals(accessController.getPrincipalId(mockCertificate), expectedId, "Expected the correct principal ID."); + assertEquals( + accessController.getPrincipalId(null), + NoOpDynamicAccessController.USER_UNKNOWN, + "Expected USER_UNKNOWN for null certificate."); + } + + @Test + public void testGetAccessControlledResources() { + Set resources = accessController.getAccessControlledResources(); + assertTrue(resources.isEmpty(), "Expected access-controlled resources to be empty."); + } + + @Test + public void testIsFailOpen() { + assertFalse(accessController.isFailOpen(), "Expected isFailOpen to return false."); + } +} diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/controllerapi/LeaderControllerResponseTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/controllerapi/LeaderControllerResponseTest.java new file mode 100644 index 00000000000..2653643b516 --- /dev/null +++ b/internal/venice-common/src/test/java/com/linkedin/venice/controllerapi/LeaderControllerResponseTest.java @@ -0,0 +1,59 @@ +package com.linkedin.venice.controllerapi; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNull; + +import org.testng.annotations.Test; + + +public class LeaderControllerResponseTest { + @Test + public void testLeaderControllerResponse() { + LeaderControllerResponse response = new LeaderControllerResponse(); + + // Testing getters and setters + String cluster = "test-cluster"; + String url = "http://leader-url"; + String secureUrl = "https://secure-leader-url"; + String grpcUrl = "grpc://leader-grpc-url"; + String secureGrpcUrl = "grpcs://secure-leader-grpc-url"; + + response.setCluster(cluster); + assertEquals(response.getCluster(), cluster, "Cluster name should match the set value"); + + response.setUrl(url); + assertEquals(response.getUrl(), url, "URL should match the set value"); + + response.setSecureUrl(secureUrl); + assertEquals(response.getSecureUrl(), secureUrl, "Secure URL should match the set value"); + + response.setGrpcUrl(grpcUrl); + assertEquals(response.getGrpcUrl(), grpcUrl, "gRPC URL should match the set value"); + + response.setSecureGrpcUrl(secureGrpcUrl); + assertEquals(response.getSecureGrpcUrl(), secureGrpcUrl, "Secure gRPC URL should match the set value"); + + // Testing default constructor + LeaderControllerResponse defaultResponse = new LeaderControllerResponse(); + assertNull(defaultResponse.getCluster(), "Cluster should be null by default"); + assertNull(defaultResponse.getUrl(), "URL should be null by default"); + assertNull(defaultResponse.getSecureUrl(), "Secure URL should be null by default"); + assertNull(defaultResponse.getGrpcUrl(), "gRPC URL should be null by default"); + assertNull(defaultResponse.getSecureGrpcUrl(), "Secure gRPC URL should be null by default"); + + // Testing combined constructor + LeaderControllerResponse combinedResponse = new LeaderControllerResponse(); + + combinedResponse.setCluster(cluster); + combinedResponse.setUrl(url); + combinedResponse.setSecureUrl(secureUrl); + combinedResponse.setGrpcUrl(grpcUrl); + combinedResponse.setSecureGrpcUrl(secureGrpcUrl); + + assertEquals(combinedResponse.getCluster(), cluster, "Cluster should match the set value"); + assertEquals(combinedResponse.getUrl(), url, "URL should match the set value"); + assertEquals(combinedResponse.getSecureUrl(), secureUrl, "Secure URL should match the set value"); + assertEquals(combinedResponse.getGrpcUrl(), grpcUrl, "gRPC URL should match the set value"); + assertEquals(combinedResponse.getSecureGrpcUrl(), secureGrpcUrl, "Secure gRPC URL should match the set value"); + } +} diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/controllerapi/RequestTopicForPushRequestTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/controllerapi/RequestTopicForPushRequestTest.java new file mode 100644 index 00000000000..90a79a1f89c --- /dev/null +++ b/internal/venice-common/src/test/java/com/linkedin/venice/controllerapi/RequestTopicForPushRequestTest.java @@ -0,0 +1,109 @@ +package com.linkedin.venice.controllerapi; + +import static com.linkedin.venice.meta.Version.PushType.BATCH; +import static org.mockito.Mockito.mock; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + +import java.security.cert.X509Certificate; +import java.util.Arrays; +import java.util.Collections; +import java.util.HashSet; +import org.testng.Assert; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class RequestTopicForPushRequestTest { + private RequestTopicForPushRequest request; + + @BeforeMethod + public void setUp() { + request = new RequestTopicForPushRequest("clusterA", "storeA", BATCH, "job123"); + } + + @Test + public void testRequestTopicForPushRequestConstructorArgs() { + assertEquals(request.getClusterName(), "clusterA"); + assertEquals(request.getStoreName(), "storeA"); + assertEquals(request.getPushType(), BATCH); + assertEquals(request.getPushJobId(), "job123"); + + // Invalid clusterName + IllegalArgumentException ex1 = Assert.expectThrows( + IllegalArgumentException.class, + () -> new RequestTopicForPushRequest("", "storeA", BATCH, "job123")); + assertEquals(ex1.getMessage(), "clusterName is required"); + + // Invalid storeName + IllegalArgumentException ex2 = Assert.expectThrows( + IllegalArgumentException.class, + () -> new RequestTopicForPushRequest("clusterA", "", BATCH, "job123")); + assertEquals(ex2.getMessage(), "storeName is required"); + + // Null pushType + IllegalArgumentException ex3 = Assert.expectThrows( + IllegalArgumentException.class, + () -> new RequestTopicForPushRequest("clusterA", "storeA", null, "job123")); + assertEquals(ex3.getMessage(), "pushType is required"); + + // Invalid pushJobId + IllegalArgumentException ex4 = Assert.expectThrows( + IllegalArgumentException.class, + () -> new RequestTopicForPushRequest("clusterA", "storeA", BATCH, "")); + assertEquals(ex4.getMessage(), "pushJobId is required"); + } + + @Test + public void testRequestTopicForPushRequestSettersAndGetters() { + request.setSendStartOfPush(true); + request.setSorted(true); + request.setWriteComputeEnabled(true); + request.setSourceGridFabric("fabricA"); + request.setRewindTimeInSecondsOverride(3600); + request.setDeferVersionSwap(true); + request.setTargetedRegions("regionA,regionB"); + request.setRepushSourceVersion(42); + request.setPartitioners("partitioner1,partitioner2"); + request.setCompressionDictionary("compressionDict"); + request.setEmergencySourceRegion("regionX"); + request.setSeparateRealTimeTopicEnabled(true); + + X509Certificate x509Certificate = mock(X509Certificate.class); + request.setCertificateInRequest(x509Certificate); + + assertTrue(request.isSendStartOfPush()); + assertTrue(request.isSorted()); + assertTrue(request.isWriteComputeEnabled()); + assertEquals(request.getSourceGridFabric(), "fabricA"); + assertEquals(request.getRewindTimeInSecondsOverride(), 3600); + assertTrue(request.isDeferVersionSwap()); + assertEquals(request.getTargetedRegions(), "regionA,regionB"); + assertEquals(request.getRepushSourceVersion(), 42); + assertEquals(request.getPartitioners(), new HashSet<>(Arrays.asList("partitioner1", "partitioner2"))); + assertEquals(request.getCompressionDictionary(), "compressionDict"); + assertEquals(request.getEmergencySourceRegion(), "regionX"); + assertTrue(request.isSeparateRealTimeTopicEnabled()); + assertEquals(request.getCertificateInRequest(), x509Certificate); + } + + @Test + public void testSetPartitionersValidAndEmptyCases() { + // Valid partitioners + request.setPartitioners("partitioner1"); + assertEquals(request.getPartitioners(), new HashSet<>(Collections.singletonList("partitioner1"))); + request.setPartitioners("partitioner1,partitioner2"); + assertEquals(request.getPartitioners(), new HashSet<>(Arrays.asList("partitioner1", "partitioner2"))); + + // Empty set + request.setPartitioners(Collections.emptySet()); + assertEquals(request.getPartitioners(), Collections.emptySet()); + + // Null and empty string + request.setPartitioners((String) null); + assertEquals(request.getPartitioners(), Collections.emptySet()); + + request.setPartitioners(""); + assertEquals(request.getPartitioners(), Collections.emptySet()); + } +} diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/controllerapi/transport/GrpcRequestResponseConverterTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/controllerapi/transport/GrpcRequestResponseConverterTest.java new file mode 100644 index 00000000000..94bf938ad0e --- /dev/null +++ b/internal/venice-common/src/test/java/com/linkedin/venice/controllerapi/transport/GrpcRequestResponseConverterTest.java @@ -0,0 +1,147 @@ +package com.linkedin.venice.controllerapi.transport; + +import static org.mockito.Mockito.argThat; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import com.google.protobuf.Any; +import com.google.rpc.Code; +import com.google.rpc.Status; +import com.linkedin.venice.client.exceptions.VeniceClientException; +import com.linkedin.venice.controllerapi.ControllerResponse; +import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo; +import com.linkedin.venice.protocols.controller.ControllerGrpcErrorType; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcErrorInfo; +import io.grpc.StatusRuntimeException; +import io.grpc.protobuf.StatusProto; +import io.grpc.stub.StreamObserver; +import org.testng.annotations.Test; + + +public class GrpcRequestResponseConverterTest { + private static final String TEST_CLUSTER = "testCluster"; + private static final String TEST_STORE = "testStore"; + + @Test + public void testGetClusterStoreGrpcInfoFromResponse() { + // Test with all fields set + ControllerResponse response = mock(ControllerResponse.class); + when(response.getCluster()).thenReturn("testCluster"); + when(response.getName()).thenReturn("testStore"); + ClusterStoreGrpcInfo grpcInfo = GrpcRequestResponseConverter.getClusterStoreGrpcInfo(response); + assertEquals(grpcInfo.getClusterName(), "testCluster"); + assertEquals(grpcInfo.getStoreName(), "testStore"); + + // Test with null fields + when(response.getCluster()).thenReturn(null); + when(response.getName()).thenReturn(null); + grpcInfo = GrpcRequestResponseConverter.getClusterStoreGrpcInfo(response); + assertEquals(grpcInfo.getClusterName(), ""); + assertEquals(grpcInfo.getStoreName(), ""); + } + + @Test + public void testSendErrorResponse() { + StreamObserver responseObserver = mock(StreamObserver.class); + + Exception e = new Exception("Test error message"); + Code errorCode = Code.INVALID_ARGUMENT; + ControllerGrpcErrorType errorType = ControllerGrpcErrorType.BAD_REQUEST; + + GrpcRequestResponseConverter.sendErrorResponse( + io.grpc.Status.Code.INVALID_ARGUMENT, + ControllerGrpcErrorType.BAD_REQUEST, + e, + TEST_CLUSTER, + TEST_STORE, + responseObserver); + + verify(responseObserver, times(1)).onError(argThat(statusRuntimeException -> { + com.google.rpc.Status status = StatusProto.fromThrowable((StatusRuntimeException) statusRuntimeException); + + VeniceControllerGrpcErrorInfo errorInfo = null; + for (Any detail: status.getDetailsList()) { + if (detail.is(VeniceControllerGrpcErrorInfo.class)) { + try { + errorInfo = detail.unpack(VeniceControllerGrpcErrorInfo.class); + break; + } catch (Exception ignored) { + } + } + } + + assertNotNull(errorInfo); + assertEquals(errorInfo.getErrorType(), errorType); + assertEquals(errorInfo.getErrorMessage(), "Test error message"); + assertEquals(errorInfo.getClusterName(), "testCluster"); + assertEquals(errorInfo.getStoreName(), "testStore"); + assertEquals(status.getCode(), errorCode.getNumber()); + + return true; + })); + } + + @Test + public void testParseControllerGrpcError() { + // Create a valid VeniceControllerGrpcErrorInfo + VeniceControllerGrpcErrorInfo errorInfo = VeniceControllerGrpcErrorInfo.newBuilder() + .setErrorType(ControllerGrpcErrorType.BAD_REQUEST) + .setErrorMessage("Invalid input") + .setStatusCode(Code.INVALID_ARGUMENT.getNumber()) + .build(); + + // Wrap in a com.google.rpc.Status + Status rpcStatus = + Status.newBuilder().setCode(Code.INVALID_ARGUMENT.getNumber()).addDetails(Any.pack(errorInfo)).build(); + + // Convert to StatusRuntimeException + StatusRuntimeException exception = StatusProto.toStatusRuntimeException(rpcStatus); + + // Parse the error + VeniceControllerGrpcErrorInfo parsedError = GrpcRequestResponseConverter.parseControllerGrpcError(exception); + + // Assert the parsed error matches the original + assertEquals(parsedError.getErrorType(), ControllerGrpcErrorType.BAD_REQUEST); + assertEquals(parsedError.getErrorMessage(), "Invalid input"); + assertEquals(parsedError.getStatusCode(), Code.INVALID_ARGUMENT.getNumber()); + } + + @Test + public void testParseControllerGrpcErrorWithNoDetails() { + // Create an exception with no details + Status rpcStatus = Status.newBuilder().setCode(Code.UNKNOWN.getNumber()).build(); + StatusRuntimeException exception = StatusProto.toStatusRuntimeException(rpcStatus); + + VeniceClientException thrownException = expectThrows( + VeniceClientException.class, + () -> GrpcRequestResponseConverter.parseControllerGrpcError(exception)); + + assertEquals(thrownException.getMessage(), "An unknown gRPC error occurred. Error code: UNKNOWN"); + } + + @Test + public void testParseControllerGrpcErrorWithUnpackFailure() { + // Create a corrupted detail + Any corruptedDetail = Any.newBuilder() + .setTypeUrl("type.googleapis.com/" + VeniceControllerGrpcErrorInfo.getDescriptor().getFullName()) + .setValue(com.google.protobuf.ByteString.copyFromUtf8("corrupted data")) + .build(); + + Status rpcStatus = + Status.newBuilder().setCode(Code.INVALID_ARGUMENT.getNumber()).addDetails(corruptedDetail).build(); + + StatusRuntimeException exception = StatusProto.toStatusRuntimeException(rpcStatus); + + VeniceClientException thrownException = expectThrows(VeniceClientException.class, () -> { + GrpcRequestResponseConverter.parseControllerGrpcError(exception); + }); + + assertTrue(thrownException.getMessage().contains("Failed to unpack error details")); + } +} diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/grpc/GrpcUtilsTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/grpc/GrpcUtilsTest.java index d89a2d36751..6059d27e94d 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/grpc/GrpcUtilsTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/grpc/GrpcUtilsTest.java @@ -1,14 +1,29 @@ package com.linkedin.venice.grpc; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; import com.linkedin.venice.acl.handler.AccessResult; +import com.linkedin.venice.client.exceptions.VeniceClientException; +import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.security.SSLFactory; import com.linkedin.venice.utils.SslUtils; +import io.grpc.Attributes; +import io.grpc.ChannelCredentials; +import io.grpc.Grpc; +import io.grpc.InsecureChannelCredentials; +import io.grpc.ServerCall; import io.grpc.Status; +import io.grpc.TlsChannelCredentials; +import java.security.cert.Certificate; +import java.security.cert.X509Certificate; import javax.net.ssl.KeyManager; +import javax.net.ssl.SSLPeerUnverifiedException; +import javax.net.ssl.SSLSession; import javax.net.ssl.TrustManager; import org.testng.annotations.BeforeTest; import org.testng.annotations.Test; @@ -70,4 +85,124 @@ public void testHttpResponseStatusToGrpcStatus() { grpcStatus.getDescription(), "Mismatch in error description for the mapped grpc status"); } + + @Test + public void testBuildChannelCredentials() { + // Case 1: sslFactory is null, expect InsecureChannelCredentials + ChannelCredentials credentials = GrpcUtils.buildChannelCredentials(null); + assertTrue( + credentials instanceof InsecureChannelCredentials, + "Expected InsecureChannelCredentials when sslFactory is null"); + + // Case 2: Valid sslFactory, expect TlsChannelCredentials + SSLFactory validSslFactory = SslUtils.getVeniceLocalSslFactory(); + credentials = GrpcUtils.buildChannelCredentials(validSslFactory); + assertTrue( + credentials instanceof TlsChannelCredentials, + "Expected TlsChannelCredentials when sslFactory is provided"); + + // Case 3: SSLFactory throws an exception when initializing credentials + SSLFactory faultySslFactory = mock(SSLFactory.class); + Exception exception = + expectThrows(VeniceClientException.class, () -> GrpcUtils.buildChannelCredentials(faultySslFactory)); + assertEquals( + exception.getMessage(), + "Failed to initialize SSL channel credentials for Venice gRPC Transport Client"); + } + + @Test + public void testExtractGrpcClientCertWithValidCertificate() throws SSLPeerUnverifiedException { + // Mock SSLSession and Certificate + SSLSession sslSession = mock(SSLSession.class); + X509Certificate x509Certificate = mock(X509Certificate.class); + + // Mock the ServerCall and its attributes + Attributes attributes = Attributes.newBuilder().set(Grpc.TRANSPORT_ATTR_SSL_SESSION, sslSession).build(); + ServerCall call = mock(ServerCall.class); + when(call.getAttributes()).thenReturn(attributes); + when(sslSession.getPeerCertificates()).thenReturn(new Certificate[] { x509Certificate }); + + // Extract the certificate + X509Certificate extractedCertificate = GrpcUtils.extractGrpcClientCert(call); + + // Verify the returned certificate + assertEquals(extractedCertificate, x509Certificate); + } + + @Test + public void testExtractGrpcClientCertWithNullSslSession() { + // Mock the ServerCall with null SSLSession + Attributes attributes = Attributes.newBuilder().set(Grpc.TRANSPORT_ATTR_SSL_SESSION, null).build(); + ServerCall call = mock(ServerCall.class); + when(call.getAttributes()).thenReturn(attributes); + when(call.getAuthority()).thenReturn("test-authority"); + + // Expect a VeniceException + VeniceException thrownException = expectThrows(VeniceException.class, () -> { + GrpcUtils.extractGrpcClientCert(call); + }); + + // Verify the exception message + assertEquals(thrownException.getMessage(), "Failed to obtain SSL session"); + } + + @Test + public void testExtractGrpcClientCertWithPeerCertificateNotX509() throws SSLPeerUnverifiedException { + // Mock SSLSession and Certificate + SSLSession sslSession = mock(SSLSession.class); + Certificate nonX509Certificate = mock(Certificate.class); + + // Mock the ServerCall and its attributes + Attributes attributes = Attributes.newBuilder().set(Grpc.TRANSPORT_ATTR_SSL_SESSION, sslSession).build(); + ServerCall call = mock(ServerCall.class); + when(call.getAttributes()).thenReturn(attributes); + when(sslSession.getPeerCertificates()).thenReturn(new Certificate[] { nonX509Certificate }); + + // Expect IllegalArgumentException + IllegalArgumentException thrownException = + expectThrows(IllegalArgumentException.class, () -> GrpcUtils.extractGrpcClientCert(call)); + + // Verify the exception message + assertTrue( + thrownException.getMessage() + .contains("Only certificates of type java.security.cert.X509Certificate are supported")); + } + + @Test + public void testExtractGrpcClientCertWithNullPeerCertificates() throws SSLPeerUnverifiedException { + // Mock SSLSession with null peer certificates + SSLSession sslSession = mock(SSLSession.class); + when(sslSession.getPeerCertificates()).thenReturn(null); + + // Mock the ServerCall and its attributes + Attributes attributes = Attributes.newBuilder().set(Grpc.TRANSPORT_ATTR_SSL_SESSION, sslSession).build(); + ServerCall call = mock(ServerCall.class); + when(call.getAttributes()).thenReturn(attributes); + + // Expect NullPointerException or VeniceException + NullPointerException thrownException = + expectThrows(NullPointerException.class, () -> GrpcUtils.extractGrpcClientCert(call)); + + // Verify the exception is thrown + assertNotNull(thrownException); + } + + @Test + public void testExtractGrpcClientCertWithEmptyPeerCertificates() throws SSLPeerUnverifiedException { + // Mock SSLSession with empty peer certificates + SSLSession sslSession = mock(SSLSession.class); + when(sslSession.getPeerCertificates()).thenReturn(new Certificate[] {}); + + // Mock the ServerCall and its attributes + Attributes attributes = Attributes.newBuilder().set(Grpc.TRANSPORT_ATTR_SSL_SESSION, sslSession).build(); + ServerCall call = mock(ServerCall.class); + when(call.getAttributes()).thenReturn(attributes); + + // Expect IndexOutOfBoundsException + IndexOutOfBoundsException thrownException = + expectThrows(IndexOutOfBoundsException.class, () -> GrpcUtils.extractGrpcClientCert(call)); + + // Verify the exception is thrown + assertNotNull(thrownException); + } } diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerConfigTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerConfigTest.java new file mode 100644 index 00000000000..82ccaa6ea0e --- /dev/null +++ b/internal/venice-common/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerConfigTest.java @@ -0,0 +1,112 @@ +package com.linkedin.venice.grpc; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import com.linkedin.venice.security.SSLFactory; +import io.grpc.BindableService; +import io.grpc.ServerCredentials; +import io.grpc.ServerInterceptor; +import java.util.Collections; +import java.util.concurrent.Executor; +import java.util.concurrent.Executors; +import java.util.concurrent.ThreadPoolExecutor; +import org.testng.annotations.Test; + + +public class VeniceGrpcServerConfigTest { + @Test + public void testBuilderWithAllFieldsSet() { + ServerCredentials credentials = mock(ServerCredentials.class); + BindableService service = mock(BindableService.class); + ServerInterceptor interceptor = mock(ServerInterceptor.class); + SSLFactory sslFactory = mock(SSLFactory.class); + Executor executor = Executors.newSingleThreadExecutor(); + + VeniceGrpcServerConfig config = new VeniceGrpcServerConfig.Builder().setPort(8080) + .setCredentials(credentials) + .setService(service) + .setInterceptor(interceptor) + .setSslFactory(sslFactory) + .setExecutor(executor) + .build(); + + assertEquals(config.getPort(), 8080); + assertEquals(config.getCredentials(), credentials); + assertEquals(config.getService(), service); + assertEquals(config.getInterceptors(), Collections.singletonList(interceptor)); + assertEquals(config.getSslFactory(), sslFactory); + assertEquals(config.getExecutor(), executor); + } + + @Test + public void testBuilderWithDefaultInterceptors() { + BindableService service = mock(BindableService.class); + + VeniceGrpcServerConfig config = + new VeniceGrpcServerConfig.Builder().setPort(8080).setService(service).setNumThreads(2).build(); + + assertTrue(config.getInterceptors().isEmpty()); + assertNotNull(config.getExecutor()); + } + + @Test + public void testBuilderWithDefaultExecutor() { + BindableService service = mock(BindableService.class); + + VeniceGrpcServerConfig config = + new VeniceGrpcServerConfig.Builder().setPort(8080).setService(service).setNumThreads(4).build(); + + assertNotNull(config.getExecutor()); + assertEquals(((ThreadPoolExecutor) config.getExecutor()).getCorePoolSize(), 4); + } + + @Test + public void testToStringMethod() { + BindableService service = mock(BindableService.class); + when(service.toString()).thenReturn("MockService"); + + VeniceGrpcServerConfig config = + new VeniceGrpcServerConfig.Builder().setPort(9090).setService(service).setNumThreads(2).build(); + + String expectedString = "VeniceGrpcServerConfig{port=9090, service=MockService}"; + assertEquals(config.toString(), expectedString); + } + + @Test + public void testBuilderValidationWithMissingPort() { + BindableService service = mock(BindableService.class); + + IllegalArgumentException exception = expectThrows(IllegalArgumentException.class, () -> { + new VeniceGrpcServerConfig.Builder().setService(service).setNumThreads(2).build(); + }); + + assertEquals(exception.getMessage(), "Port value is required to create the gRPC server but was not provided."); + } + + @Test + public void testBuilderValidationWithMissingService() { + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new VeniceGrpcServerConfig.Builder().setPort(8080).setNumThreads(2).build()); + + assertEquals(exception.getMessage(), "A non-null gRPC service instance is required to create the server."); + } + + @Test + public void testBuilderValidationWithInvalidThreadsAndExecutor() { + BindableService service = mock(BindableService.class); + + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new VeniceGrpcServerConfig.Builder().setPort(8080).setService(service).build()); + + assertEquals( + exception.getMessage(), + "gRPC server creation requires a valid number of threads (numThreads > 0) or a non-null executor."); + } +} diff --git a/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerTest.java similarity index 80% rename from services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerTest.java rename to internal/venice-common/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerTest.java index 47a6ef204b6..3d6efb8b331 100644 --- a/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerTest.java @@ -1,13 +1,12 @@ package com.linkedin.venice.grpc; -import static org.mockito.Mockito.mock; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertTrue; import com.linkedin.venice.exceptions.VeniceException; -import com.linkedin.venice.listener.grpc.VeniceReadServiceImpl; -import com.linkedin.venice.listener.grpc.handlers.VeniceServerGrpcRequestProcessor; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcServiceGrpc; import com.linkedin.venice.security.SSLFactory; import com.linkedin.venice.utils.SslUtils; import com.linkedin.venice.utils.TestUtils; @@ -19,16 +18,15 @@ public class VeniceGrpcServerTest { + private static final int NUM_THREADS = 2; private VeniceGrpcServer grpcServer; private VeniceGrpcServerConfig.Builder serverConfig; - private VeniceServerGrpcRequestProcessor grpcRequestProcessor; @BeforeMethod void setUp() { - grpcRequestProcessor = mock(VeniceServerGrpcRequestProcessor.class); serverConfig = new VeniceGrpcServerConfig.Builder().setPort(TestUtils.getFreePort()) - .setNumThreads(10) - .setService(new VeniceReadServiceImpl(grpcRequestProcessor)); + .setService(new VeniceControllerGrpcServiceTestImpl()) + .setNumThreads(NUM_THREADS); } @Test @@ -36,10 +34,12 @@ void startServerSuccessfully() { grpcServer = new VeniceGrpcServer(serverConfig.build()); grpcServer.start(); + assertNotNull(grpcServer.getServer()); + assertTrue(grpcServer.isRunning()); assertFalse(grpcServer.isTerminated()); grpcServer.stop(); - assertTrue(grpcServer.isShutdown()); + assertFalse(grpcServer.isRunning()); } @Test @@ -65,7 +65,7 @@ void testServerShutdown() throws InterruptedException { Thread.sleep(500); grpcServer.stop(); - assertTrue(grpcServer.isShutdown()); + assertFalse(grpcServer.isRunning()); Thread.sleep(500); @@ -89,4 +89,8 @@ void testServerWithSSL() { assertFalse(serverCredentials instanceof TlsServerCredentials); assertTrue(serverCredentials instanceof InsecureServerCredentials); } + + public static class VeniceControllerGrpcServiceTestImpl + extends VeniceControllerGrpcServiceGrpc.VeniceControllerGrpcServiceImplBase { + } } diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/helix/HelixCustomizedViewOfflinePushRepositoryTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/helix/HelixCustomizedViewOfflinePushRepositoryTest.java index 7ee08c747ff..152032f6868 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/helix/HelixCustomizedViewOfflinePushRepositoryTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/helix/HelixCustomizedViewOfflinePushRepositoryTest.java @@ -1,7 +1,12 @@ package com.linkedin.venice.helix; -import static org.mockito.ArgumentMatchers.*; -import static org.mockito.Mockito.*; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import com.linkedin.venice.meta.ReadOnlyStoreRepository; import com.linkedin.venice.meta.Store; diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/helix/TestHelixLiveInstanceMonitor.java b/internal/venice-common/src/test/java/com/linkedin/venice/helix/TestHelixLiveInstanceMonitor.java index a59add061cb..7710a9c8c26 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/helix/TestHelixLiveInstanceMonitor.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/helix/TestHelixLiveInstanceMonitor.java @@ -1,7 +1,9 @@ package com.linkedin.venice.helix; -import static org.mockito.ArgumentMatchers.*; -import static org.mockito.Mockito.*; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; import org.apache.helix.zookeeper.impl.client.ZkClient; import org.mockito.Mockito; diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/meta/NameRepositoryTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/meta/NameRepositoryTest.java new file mode 100644 index 00000000000..859498e346b --- /dev/null +++ b/internal/venice-common/src/test/java/com/linkedin/venice/meta/NameRepositoryTest.java @@ -0,0 +1,76 @@ +package com.linkedin.venice.meta; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNotSame; +import static org.testng.Assert.assertSame; +import static org.testng.Assert.assertThrows; + +import org.testng.annotations.Test; + + +public class NameRepositoryTest { + @Test + public void test() { + NameRepository nameRepository = new NameRepository(); + + // Only valid stores and store-versions should be possible + assertThrows(IllegalArgumentException.class, () -> nameRepository.getStoreName("invalid--store--name")); + assertThrows(IllegalArgumentException.class, () -> nameRepository.getStoreVersionName("missingTheUnderscoreV23")); + + assertThrows(IllegalArgumentException.class, () -> new StoreName("invalid--store--name")); + assertThrows( + IllegalArgumentException.class, + () -> new StoreVersionName("missingTheUnderscoreV23", new StoreName("store"))); + + // Basic "happy path" stuff + String storeNameString = "store"; + StoreName storeName = nameRepository.getStoreName(storeNameString); + assertNotNull(storeName); + assertEquals(storeName.getName(), storeNameString); + + int versionNumber = 1; + StoreVersionName storeVersionName = nameRepository.getStoreVersionName(storeNameString, versionNumber); + assertNotNull(storeVersionName); + String storeVersionNameString = storeNameString + "_v" + versionNumber; + assertEquals(storeVersionName.getName(), storeVersionNameString); + assertEquals(storeVersionName.getStore(), storeName); + assertSame(storeVersionName.getStore(), storeName, "These should not only be equal, but also the same ref!"); + assertEquals(storeVersionName.getStoreName(), storeName.getName()); + assertEquals(storeVersionName.getVersionNumber(), versionNumber); + + StoreVersionName storeVersionName2 = nameRepository.getStoreVersionName(storeVersionNameString); + assertEquals(storeVersionName, storeVersionName2); + assertSame(storeVersionName, storeVersionName2, "These should not only be equal, but also the same ref!"); + + // Ensure that equals and hashCode works correctly, + // even if different refs are compared (which could happen if the cache grows too big) + StoreName storeNameOtherRef = new StoreName(storeNameString); + assertEquals(storeName, storeNameOtherRef); + assertEquals(storeName.toString(), storeNameOtherRef.toString()); + assertNotSame(storeName, storeNameOtherRef, "These should not be the same ref!"); + assertEquals(storeName.hashCode(), storeNameOtherRef.hashCode()); + + StoreVersionName storeVersionNameOtherRef = new StoreVersionName(storeVersionNameString, storeName); + assertEquals(storeVersionName, storeVersionNameOtherRef); + assertEquals(storeVersionName.toString(), storeVersionNameOtherRef.toString()); + assertNotSame(storeVersionName, storeVersionNameOtherRef, "These should not be the same ref!"); + assertEquals(storeVersionName.hashCode(), storeVersionNameOtherRef.hashCode()); + + // Not equals + StoreName differentStoreName = nameRepository.getStoreName("differentStore"); + assertNotNull(differentStoreName); + assertNotEquals(storeName, differentStoreName); + + StoreVersionName differentStoreVersionName = nameRepository.getStoreVersionName("differentStore", 1); + assertNotNull(differentStoreVersionName); + assertNotEquals(storeVersionName, differentStoreVersionName); + + // Nonsensical equals + assertNotEquals(storeName, new Object()); + assertNotEquals(new Object(), storeName); + assertNotEquals(storeVersionName, new Object()); + assertNotEquals(new Object(), storeVersionName); + } +} diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/meta/RetryManagerTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/meta/RetryManagerTest.java index 5f6b3426ee2..ae615b6abc5 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/meta/RetryManagerTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/meta/RetryManagerTest.java @@ -1,6 +1,7 @@ package com.linkedin.venice.meta; -import static org.mockito.Mockito.*; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; import com.linkedin.venice.utils.TestUtils; import io.tehuti.metrics.MetricsRepository; diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestInstance.java b/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestInstance.java index 47c27598430..6395054ec15 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestInstance.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestInstance.java @@ -23,4 +23,17 @@ public void parsesNodeId() { Assert.assertEquals(host.getHost(), "localhost"); Assert.assertEquals(host.getPort(), 1234); } + + @Test + public void testInstanceWithGrpcAddress() { + Instance nonGrpcInstance = new Instance("localhost_1234", "localhost", 1234); + Assert.assertEquals(nonGrpcInstance.getGrpcSslPort(), -1); + Assert.assertEquals(nonGrpcInstance.getGrpcPort(), -1); + + Instance grpcInstance = new Instance("localhost_1234", "localhost", 1234, 1235, 1236); + Assert.assertEquals(grpcInstance.getGrpcPort(), 1235); + Assert.assertEquals(grpcInstance.getGrpcSslPort(), 1236); + Assert.assertEquals(grpcInstance.getGrpcUrl(), "localhost:1235"); + Assert.assertEquals(grpcInstance.getGrpcSslUrl(), "localhost:1236"); + } } diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestVersion.java b/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestVersion.java index f28aebb2c84..a03824a31af 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestVersion.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestVersion.java @@ -1,9 +1,14 @@ package com.linkedin.venice.meta; import static com.linkedin.venice.meta.Version.VENICE_TTL_RE_PUSH_PUSH_ID_PREFIX; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; import com.fasterxml.jackson.databind.ObjectMapper; import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.meta.Version.PushType; import com.linkedin.venice.utils.ObjectMapperFactory; import com.linkedin.venice.utils.Utils; import java.io.IOException; @@ -27,7 +32,7 @@ public class TestVersion { @Test public void identifiesValidTopicNames() { String goodTopic = "my_very_good_store_v4"; - Assert.assertTrue( + assertTrue( Version.isVersionTopicOrStreamReprocessingTopic(goodTopic), goodTopic + " should parse as a valid store-version topic"); @@ -43,7 +48,7 @@ public void serializes() throws IOException { int versionNumber = 17; Version version = new VersionImpl(storeName, versionNumber); String serialized = OBJECT_MAPPER.writeValueAsString(version); - Assert.assertTrue(serialized.contains(storeName)); + assertTrue(serialized.contains(storeName)); } /** @@ -54,21 +59,21 @@ public void serializes() throws IOException { @Test public void deserializeWithWrongFields() throws IOException { Version oldParsedVersion = OBJECT_MAPPER.readValue(OLD_SERIALIZED, Version.class); - Assert.assertEquals(oldParsedVersion.getStoreName(), "store-1492637190910-78714331"); + assertEquals(oldParsedVersion.getStoreName(), "store-1492637190910-78714331"); Version newParsedVersion = OBJECT_MAPPER.readValue(EXTRA_FIELD_SERIALIZED, Version.class); - Assert.assertEquals(newParsedVersion.getStoreName(), "store-1492637190910-12345678"); + assertEquals(newParsedVersion.getStoreName(), "store-1492637190910-12345678"); Version legacyParsedVersion = OBJECT_MAPPER.readValue(MISSING_FIELD_SERIALIZED, Version.class); - Assert.assertEquals(legacyParsedVersion.getStoreName(), "store-missing"); - Assert.assertNotNull(legacyParsedVersion.getPushJobId()); // missing final field can still deserialize, just gets - // arbitrary value from constructor + assertEquals(legacyParsedVersion.getStoreName(), "store-missing"); + assertNotNull(legacyParsedVersion.getPushJobId()); // missing final field can still deserialize, just gets + // arbitrary value from constructor } @Test public void testParseStoreFromRealTimeTopic() { String validRealTimeTopic = "abc_rt"; - Assert.assertEquals(Version.parseStoreFromRealTimeTopic(validRealTimeTopic), "abc"); + assertEquals(Version.parseStoreFromRealTimeTopic(validRealTimeTopic), "abc"); String invalidRealTimeTopic = "abc"; try { Version.parseStoreFromRealTimeTopic(invalidRealTimeTopic); @@ -82,19 +87,19 @@ public void testParseStoreFromRealTimeTopic() { public void testIsTopic() { String topic = "abc_rt"; Assert.assertFalse(Version.isVersionTopic(topic)); - Assert.assertTrue(Version.isRealTimeTopic(topic)); + assertTrue(Version.isRealTimeTopic(topic)); topic = "abc"; Assert.assertFalse(Version.isVersionTopic(topic)); topic = "abc_v12df"; Assert.assertFalse(Version.isVersionTopic(topic)); topic = "abc_v123"; - Assert.assertTrue(Version.isVersionTopic(topic)); + assertTrue(Version.isVersionTopic(topic)); Assert.assertFalse(Version.isRealTimeTopic(topic)); - Assert.assertTrue(Version.isVersionTopicOrStreamReprocessingTopic(topic)); + assertTrue(Version.isVersionTopicOrStreamReprocessingTopic(topic)); topic = "abc_v123_sr"; Assert.assertFalse(Version.isVersionTopic(topic)); - Assert.assertTrue(Version.isStreamReprocessingTopic(topic)); - Assert.assertTrue(Version.isVersionTopicOrStreamReprocessingTopic(topic)); + assertTrue(Version.isStreamReprocessingTopic(topic)); + assertTrue(Version.isVersionTopicOrStreamReprocessingTopic(topic)); topic = "abc_v12ab3_sr"; Assert.assertFalse(Version.isVersionTopic(topic)); Assert.assertFalse(Version.isStreamReprocessingTopic(topic)); @@ -110,43 +115,101 @@ public void testIsATopicThatIsVersioned() { String topic = "abc_rt"; Assert.assertFalse(Version.isATopicThatIsVersioned(topic)); topic = "abc_v1_sr"; - Assert.assertTrue(Version.isATopicThatIsVersioned(topic)); + assertTrue(Version.isATopicThatIsVersioned(topic)); topic = "abc_v1"; - Assert.assertTrue(Version.isATopicThatIsVersioned(topic)); + assertTrue(Version.isATopicThatIsVersioned(topic)); topic = "abc_v1_cc"; - Assert.assertTrue(Version.isATopicThatIsVersioned(topic)); + assertTrue(Version.isATopicThatIsVersioned(topic)); String pushId = VENICE_TTL_RE_PUSH_PUSH_ID_PREFIX + System.currentTimeMillis(); - Assert.assertTrue(Version.isPushIdTTLRePush(pushId)); + assertTrue(Version.isPushIdTTLRePush(pushId)); } @Test public void testParseStoreFromKafkaTopicName() { String storeName = "abc"; String topic = "abc_rt"; - Assert.assertEquals(Version.parseStoreFromKafkaTopicName(topic), storeName); + assertEquals(Version.parseStoreFromKafkaTopicName(topic), storeName); topic = "abc_v1"; - Assert.assertEquals(Version.parseStoreFromKafkaTopicName(topic), storeName); + assertEquals(Version.parseStoreFromKafkaTopicName(topic), storeName); topic = "abc_v1_cc"; - Assert.assertEquals(Version.parseStoreFromKafkaTopicName(topic), storeName); + assertEquals(Version.parseStoreFromKafkaTopicName(topic), storeName); } @Test public void testParseVersionFromKafkaTopicName() { int version = 1; String topic = "abc_v1"; - Assert.assertEquals(Version.parseVersionFromVersionTopicName(topic), version); + assertEquals(Version.parseVersionFromVersionTopicName(topic), version); topic = "abc_v1_cc"; - Assert.assertEquals(Version.parseVersionFromKafkaTopicName(topic), version); + assertEquals(Version.parseVersionFromKafkaTopicName(topic), version); } @Test void testVersionStatus() { for (VersionStatus status: VersionStatus.values()) { if (status == VersionStatus.KILLED) { - Assert.assertTrue(VersionStatus.isVersionKilled(status)); + assertTrue(VersionStatus.isVersionKilled(status)); } else { Assert.assertFalse(VersionStatus.isVersionKilled(status)); } } } + + @Test + public void testExtractPushType() { + // Case 1: Valid push types + assertEquals(PushType.extractPushType("BATCH"), PushType.BATCH); + assertEquals(PushType.extractPushType("STREAM_REPROCESSING"), PushType.STREAM_REPROCESSING); + assertEquals(PushType.extractPushType("STREAM"), PushType.STREAM); + assertEquals(PushType.extractPushType("INCREMENTAL"), PushType.INCREMENTAL); + + // Case 2: Invalid push type + String invalidType = "INVALID_TYPE"; + IllegalArgumentException invalidException = + expectThrows(IllegalArgumentException.class, () -> PushType.extractPushType(invalidType)); + assertTrue(invalidException.getMessage().contains(invalidType)); + assertTrue(invalidException.getMessage().contains("Valid push types are")); + + // Case 3: Case sensitivity + String lowerCaseType = "batch"; + IllegalArgumentException caseException = + expectThrows(IllegalArgumentException.class, () -> PushType.extractPushType(lowerCaseType)); + assertTrue(caseException.getMessage().contains(lowerCaseType)); + + // Case 4: Empty string + String emptyInput = ""; + IllegalArgumentException emptyException = + expectThrows(IllegalArgumentException.class, () -> PushType.extractPushType(emptyInput)); + assertTrue(emptyException.getMessage().contains(emptyInput)); + + // Case 5: Null input + IllegalArgumentException exception = + expectThrows(IllegalArgumentException.class, () -> PushType.extractPushType(null)); + assertNotNull(exception); + } + + @Test + public void testValueOfIntReturnsPushType() { + // Case 1: Valid integer values + assertEquals(PushType.valueOf(0), PushType.BATCH); + assertEquals(PushType.valueOf(1), PushType.STREAM_REPROCESSING); + assertEquals(PushType.valueOf(2), PushType.STREAM); + assertEquals(PushType.valueOf(3), PushType.INCREMENTAL); + + // Case 2: Invalid integer value (negative) + int invalidNegative = -1; + VeniceException negativeException = expectThrows(VeniceException.class, () -> PushType.valueOf(invalidNegative)); + assertTrue(negativeException.getMessage().contains("Invalid push type with int value: " + invalidNegative)); + + // Case 3: Invalid integer value (positive out of range) + int invalidPositive = 999; + VeniceException positiveException = expectThrows(VeniceException.class, () -> PushType.valueOf(invalidPositive)); + assertTrue(positiveException.getMessage().contains("Invalid push type with int value: " + invalidPositive)); + + // Case 4: Edge case - Valid minimum value + assertEquals(PushType.valueOf(0), PushType.BATCH); + + // Case 5: Edge case - Valid maximum value + assertEquals(PushType.valueOf(3), PushType.INCREMENTAL); + } } diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestZKStore.java b/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestZKStore.java index 5d4283d1ee7..509380ef1a9 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestZKStore.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/meta/TestZKStore.java @@ -375,10 +375,10 @@ public void testValidStoreNames() { List valid = Arrays.asList("foo", "Bar", "foo_bar", "foo-bar", "f00Bar"); List invalid = Arrays.asList("foo bar", "foo.bar", " foo", ".bar", "!", "@", "#", "$", "%"); for (String name: valid) { - Assert.assertTrue(Store.isValidStoreName(name)); + Assert.assertTrue(StoreName.isValidStoreName(name)); } for (String name: invalid) { - Assert.assertFalse(Store.isValidStoreName(name)); + Assert.assertFalse(StoreName.isValidStoreName(name)); } } diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/PubSubTopicConfigurationTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/PubSubTopicConfigurationTest.java new file mode 100644 index 00000000000..75a2a8affad --- /dev/null +++ b/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/PubSubTopicConfigurationTest.java @@ -0,0 +1,121 @@ +package com.linkedin.venice.pubsub; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotEquals; +import static org.testng.Assert.assertNotSame; +import static org.testng.Assert.assertNull; + +import java.util.Optional; +import org.testng.annotations.Test; + + +public class PubSubTopicConfigurationTest { + @Test + public void testPubSubTopicConfiguration() throws CloneNotSupportedException { + // Case 1: Verify construction and getters + Optional retentionInMs = Optional.of(3600000L); + boolean isLogCompacted = true; + Optional minInSyncReplicas = Optional.of(2); + Long minLogCompactionLagMs = 60000L; + Optional maxLogCompactionLagMs = Optional.of(120000L); + + PubSubTopicConfiguration config = new PubSubTopicConfiguration( + retentionInMs, + isLogCompacted, + minInSyncReplicas, + minLogCompactionLagMs, + maxLogCompactionLagMs); + + assertEquals(config.retentionInMs(), retentionInMs, "Retention in ms should match the provided value."); + assertEquals(config.isLogCompacted(), isLogCompacted, "Log compaction flag should match the provided value."); + assertEquals( + config.minInSyncReplicas(), + minInSyncReplicas, + "Min in-sync replicas should match the provided value."); + assertEquals( + config.minLogCompactionLagMs(), + minLogCompactionLagMs, + "Min log compaction lag ms should match the provided value."); + assertEquals( + config.getMaxLogCompactionLagMs(), + maxLogCompactionLagMs, + "Max log compaction lag ms should match the provided value."); + + // Case 2: Verify setters + config.setLogCompacted(false); + config.setRetentionInMs(Optional.of(7200000L)); + config.setMinInSyncReplicas(Optional.of(3)); + config.setMinLogCompactionLagMs(120000L); + config.setMaxLogCompactionLagMs(Optional.of(180000L)); + + assertFalse(config.isLogCompacted(), "Log compaction flag should be updated."); + assertEquals(config.retentionInMs(), Optional.of(7200000L), "Retention in ms should be updated."); + assertEquals(config.minInSyncReplicas(), Optional.of(3), "Min in-sync replicas should be updated."); + assertEquals(config.minLogCompactionLagMs(), Long.valueOf(120000L), "Min log compaction lag ms should be updated."); + assertEquals( + config.getMaxLogCompactionLagMs(), + Optional.of(180000L), + "Max log compaction lag ms should be updated."); + + // Case 3: Verify cloning + PubSubTopicConfiguration clonedConfig = config.clone(); + + assertNotSame(clonedConfig, config, "Cloned object should not be the same as the original."); + assertEquals( + clonedConfig.retentionInMs(), + config.retentionInMs(), + "Retention in ms should be identical in the cloned object."); + assertEquals( + clonedConfig.isLogCompacted(), + config.isLogCompacted(), + "Log compaction flag should be identical in the cloned object."); + assertEquals( + clonedConfig.minInSyncReplicas(), + config.minInSyncReplicas(), + "Min in-sync replicas should be identical in the cloned object."); + assertEquals( + clonedConfig.minLogCompactionLagMs(), + config.minLogCompactionLagMs(), + "Min log compaction lag ms should be identical in the cloned object."); + assertEquals( + clonedConfig.getMaxLogCompactionLagMs(), + config.getMaxLogCompactionLagMs(), + "Max log compaction lag ms should be identical in the cloned object."); + + clonedConfig.setLogCompacted(true); + clonedConfig.setRetentionInMs(Optional.of(14400000L)); + clonedConfig.setMinInSyncReplicas(Optional.of(4)); + clonedConfig.setMinLogCompactionLagMs(180000L); + + assertNotEquals( + config.isLogCompacted(), + clonedConfig.isLogCompacted(), + "Log compaction flag should be different in the cloned object."); + assertNotEquals( + config.retentionInMs(), + clonedConfig.retentionInMs(), + "Retention in ms should be different in the cloned object."); + assertNotEquals( + config.minInSyncReplicas(), + clonedConfig.minInSyncReplicas(), + "Min in-sync replicas should be different in the cloned object."); + assertNotEquals( + config.minLogCompactionLagMs(), + clonedConfig.minLogCompactionLagMs(), + "Min log compaction lag ms should be different in the cloned object."); + assertEquals( + config.getMaxLogCompactionLagMs(), + clonedConfig.getMaxLogCompactionLagMs(), + "Max log compaction lag ms should be identical in the cloned object."); + + // Case 4: Verify edge cases + PubSubTopicConfiguration emptyConfig = + new PubSubTopicConfiguration(Optional.empty(), false, Optional.empty(), null, Optional.empty()); + + assertFalse(emptyConfig.retentionInMs().isPresent(), "Retention in ms should be empty."); + assertFalse(emptyConfig.minInSyncReplicas().isPresent(), "Min in-sync replicas should be empty."); + assertNull(emptyConfig.minLogCompactionLagMs(), "Min log compaction lag ms should be null."); + assertFalse(emptyConfig.getMaxLogCompactionLagMs().isPresent(), "Max log compaction lag ms should be empty."); + } +} diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/adapter/kafka/admin/ApacheKafkaAdminConfigTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/adapter/kafka/admin/ApacheKafkaAdminConfigTest.java index 1f2a998d7ba..6058c5cf5ce 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/adapter/kafka/admin/ApacheKafkaAdminConfigTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/adapter/kafka/admin/ApacheKafkaAdminConfigTest.java @@ -1,6 +1,6 @@ package com.linkedin.venice.pubsub.adapter.kafka.admin; -import static org.testng.Assert.*; +import static org.testng.Assert.assertEquals; import com.linkedin.venice.pubsub.api.PubSubSecurityProtocol; import com.linkedin.venice.utils.VeniceProperties; diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/manager/TopicManagerTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/manager/TopicManagerTest.java index 346a0476017..95b2f5f1fef 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/manager/TopicManagerTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/manager/TopicManagerTest.java @@ -4,11 +4,13 @@ import static org.mockito.ArgumentMatchers.anyBoolean; import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.any; +import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.times; +import static org.mockito.Mockito.when; import com.github.benmanes.caffeine.cache.Cache; import com.linkedin.venice.exceptions.VeniceException; @@ -552,6 +554,51 @@ public void testGetAndUpdateTopicRetentionForNonExistingTopic() { () -> topicManager.updateTopicRetention(nonExistingTopic, TimeUnit.DAYS.toMillis(1))); } + @Test + public void testGetAndUpdateTopicRetentionWithRetriesForNonExistingTopic() { + PubSubTopic existingTopic = pubSubTopicRepository.getTopic(TestUtils.getUniqueTopicString("existing-topic")); + PubSubTopic nonExistentTopic = pubSubTopicRepository.getTopic(TestUtils.getUniqueTopicString("non-existing-topic")); + PubSubAdminAdapter mockPubSubAdminAdapter = mock(PubSubAdminAdapter.class); + PubSubAdminAdapterFactory adminAdapterFactory = mock(PubSubAdminAdapterFactory.class); + PubSubConsumerAdapterFactory consumerAdapterFactory = mock(PubSubConsumerAdapterFactory.class); + PubSubConsumerAdapter mockPubSubConsumer = mock(PubSubConsumerAdapter.class); + doReturn(mockPubSubConsumer).when(consumerAdapterFactory).create(any(), anyBoolean(), any(), anyString()); + doReturn(mockPubSubAdminAdapter).when(adminAdapterFactory).create(any(), eq(pubSubTopicRepository)); + + PubSubTopicConfiguration topicProperties = + new PubSubTopicConfiguration(Optional.of(TimeUnit.DAYS.toMillis(1)), true, Optional.of(1), 4L, Optional.of(5L)); + when(mockPubSubAdminAdapter.getTopicConfigWithRetry(existingTopic)).thenReturn(topicProperties); + when(mockPubSubAdminAdapter.getTopicConfigWithRetry(nonExistentTopic)).thenReturn(topicProperties); + + doNothing().when(mockPubSubAdminAdapter).setTopicConfig(eq(existingTopic), any(PubSubTopicConfiguration.class)); + doThrow(new PubSubTopicDoesNotExistException("Topic does not exist")).when(mockPubSubAdminAdapter) + .setTopicConfig(eq(nonExistentTopic), any(PubSubTopicConfiguration.class)); + + TopicManagerContext topicManagerContext = + new TopicManagerContext.Builder().setPubSubPropertiesSupplier(k -> VeniceProperties.empty()) + .setPubSubTopicRepository(pubSubTopicRepository) + .setPubSubAdminAdapterFactory(adminAdapterFactory) + .setPubSubConsumerAdapterFactory(consumerAdapterFactory) + .setTopicDeletionStatusPollIntervalMs(100) + .setTopicMetadataFetcherConsumerPoolSize(1) + .setTopicMetadataFetcherThreadPoolSize(1) + .setTopicMinLogCompactionLagMs(MIN_COMPACTION_LAG) + .build(); + + try (TopicManager topicManagerForThisTest = + new TopicManagerRepository(topicManagerContext, "test").getLocalTopicManager()) { + Assert.assertFalse( + topicManagerForThisTest.updateTopicRetentionWithRetries(existingTopic, TimeUnit.DAYS.toMillis(1)), + "Topic should not be updated since it already has the same retention"); + Assert.assertTrue( + topicManagerForThisTest.updateTopicRetentionWithRetries(existingTopic, TimeUnit.DAYS.toMillis(5)), + "Topic should be updated since it has different retention"); + Assert.assertThrows( + PubSubClientRetriableException.class, + () -> topicManagerForThisTest.updateTopicRetentionWithRetries(nonExistentTopic, TimeUnit.DAYS.toMillis(2))); + } + } + @Test public void testUpdateTopicCompactionPolicyForNonExistingTopic() { PubSubTopic nonExistingTopic = pubSubTopicRepository.getTopic(TestUtils.getUniqueTopicString("non-existing-topic")); diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/manager/TopicMetadataFetcherTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/manager/TopicMetadataFetcherTest.java index bca7d156296..e2f3ff2ad9e 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/manager/TopicMetadataFetcherTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/pubsub/manager/TopicMetadataFetcherTest.java @@ -29,6 +29,7 @@ import com.linkedin.venice.kafka.protocol.ProducerMetadata; import com.linkedin.venice.message.KafkaKey; import com.linkedin.venice.pubsub.ImmutablePubSubMessage; +import com.linkedin.venice.pubsub.PubSubConstants; import com.linkedin.venice.pubsub.PubSubTopicPartitionImpl; import com.linkedin.venice.pubsub.PubSubTopicPartitionInfo; import com.linkedin.venice.pubsub.PubSubTopicRepository; @@ -262,7 +263,7 @@ public void testGetTopicLatestOffsets() { assertEquals(res.get(1), 222L); assertEquals( topicMetadataFetcher.getLatestOffsetCachedNonBlocking(new PubSubTopicPartitionImpl(pubSubTopic, 0)), - -1); + PubSubConstants.UNKNOWN_LATEST_OFFSET); verify(consumerMock, times(3)).partitionsFor(pubSubTopic); verify(consumerMock, times(1)).endOffsets(eq(offsetsMap.keySet()), any(Duration.class)); diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/pushstatushelper/PushStatusStoreReaderTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/pushstatushelper/PushStatusStoreReaderTest.java index 888ad175aa2..1c25cebaa45 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/pushstatushelper/PushStatusStoreReaderTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/pushstatushelper/PushStatusStoreReaderTest.java @@ -3,7 +3,7 @@ import static com.linkedin.venice.common.PushStatusStoreUtils.SERVER_INCREMENTAL_PUSH_PREFIX; import static com.linkedin.venice.common.PushStatusStoreUtils.getServerIncrementalPushKey; import static com.linkedin.venice.pushmonitor.ExecutionStatus.END_OF_INCREMENTAL_PUSH_RECEIVED; -import static org.mockito.ArgumentMatchers.*; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.any; import static org.mockito.Mockito.anySet; import static org.mockito.Mockito.doCallRealMethod; diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/serialization/avro/AvroSpecificStoreDeserializerCacheTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/serialization/avro/AvroSpecificStoreDeserializerCacheTest.java index a59673b8ca0..2518553090a 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/serialization/avro/AvroSpecificStoreDeserializerCacheTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/serialization/avro/AvroSpecificStoreDeserializerCacheTest.java @@ -1,6 +1,11 @@ package com.linkedin.venice.serialization.avro; -import static org.mockito.Mockito.*; +import static org.mockito.Mockito.anyInt; +import static org.mockito.Mockito.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertSame; diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/utils/KafkaSSLUtilsTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/utils/KafkaSSLUtilsTest.java index ef67ef05aee..03bbfb87914 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/utils/KafkaSSLUtilsTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/utils/KafkaSSLUtilsTest.java @@ -1,6 +1,7 @@ package com.linkedin.venice.utils; -import static org.testng.Assert.*; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; import com.linkedin.venice.pubsub.api.PubSubSecurityProtocol; import org.testng.annotations.Test; diff --git a/internal/venice-common/src/test/java/com/linkedin/venice/utils/UtilsTest.java b/internal/venice-common/src/test/java/com/linkedin/venice/utils/UtilsTest.java index ca0abdd4aa0..6a6280295d4 100644 --- a/internal/venice-common/src/test/java/com/linkedin/venice/utils/UtilsTest.java +++ b/internal/venice-common/src/test/java/com/linkedin/venice/utils/UtilsTest.java @@ -1,6 +1,7 @@ package com.linkedin.venice.utils; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertThrows; @@ -9,6 +10,7 @@ import static org.testng.Assert.fail; import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.exceptions.VeniceHttpException; import com.linkedin.venice.meta.HybridStoreConfig; import com.linkedin.venice.meta.Store; import com.linkedin.venice.meta.StoreInfo; @@ -28,7 +30,9 @@ import java.util.Map; import java.util.Set; import java.util.TreeMap; +import org.apache.http.HttpStatus; import org.testng.Assert; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; import org.testng.collections.Lists; @@ -256,7 +260,7 @@ void testGetRealTimeTopicNameWithStore() { when(mockHybridConfig.getRealTimeTopicName()).thenReturn("RealTimeTopic"); String result = Utils.getRealTimeTopicName(mockStore); - assertEquals("RealTimeTopic", result); + assertEquals(result, "RealTimeTopic"); } @Test @@ -273,7 +277,7 @@ void testGetRealTimeTopicNameWithStoreInfo() { when(mockHybridConfig.getRealTimeTopicName()).thenReturn("RealTimeTopic"); String result = Utils.getRealTimeTopicName(mockStoreInfo); - assertEquals("RealTimeTopic", result); + assertEquals(result, "RealTimeTopic"); } @Test @@ -283,13 +287,13 @@ void testGetRealTimeTopicNameWithHybridConfig() { when(mockHybridConfig.getRealTimeTopicName()).thenReturn("RealTimeTopic"); String result = Utils.getRealTimeTopicName("TestStore", Collections.EMPTY_LIST, 1, mockHybridConfig); - assertEquals("RealTimeTopic", result); + assertEquals(result, "RealTimeTopic"); } @Test void testGetRealTimeTopicNameWithoutHybridConfig() { String result = Utils.getRealTimeTopicName("TestStore", Collections.EMPTY_LIST, 0, null); - assertEquals("TestStore" + Version.REAL_TIME_TOPIC_SUFFIX, result); + assertEquals(result, "TestStore" + Version.REAL_TIME_TOPIC_SUFFIX); } @Test @@ -321,7 +325,7 @@ void testGetRealTimeTopicNameWithExceptionHandling() { when(mockVersion2.isHybrid()).thenReturn(false); String result = Utils.getRealTimeTopicName("TestStore", Lists.newArrayList(mockVersion1, mockVersion2), 1, null); - assertEquals("TestStore" + Version.REAL_TIME_TOPIC_SUFFIX, result); + assertEquals(result, "TestStore" + Version.REAL_TIME_TOPIC_SUFFIX); } @Test @@ -334,7 +338,7 @@ void testGetRealTimeTopicNameWithVersion() { when(mockHybridConfig.getRealTimeTopicName()).thenReturn("RealTimeTopic"); String result = Utils.getRealTimeTopicName(mockVersion); - assertEquals("RealTimeTopic", result); + assertEquals(result, "RealTimeTopic"); } @Test @@ -346,7 +350,60 @@ void testGetRealTimeTopicNameWithNonHybridVersion() { when(mockVersion.getHybridStoreConfig()).thenReturn(null); when(mockVersion.getStoreName()).thenReturn("TestStore"); String result = Utils.getRealTimeTopicName(mockVersion); - assertEquals("TestStore" + Version.REAL_TIME_TOPIC_SUFFIX, result); + assertEquals(result, "TestStore" + Version.REAL_TIME_TOPIC_SUFFIX); + } + + @Test + void testRealTimeTopicNameWithHybridConfig() { + // Mock the Store and HybridStoreConfig + Store store = mock(Store.class); + HybridStoreConfig hybridStoreConfig = mock(HybridStoreConfig.class); + + // Define behavior + when(store.getHybridStoreConfig()).thenReturn(hybridStoreConfig); + when(store.getName()).thenReturn("test-store"); + when(hybridStoreConfig.getRealTimeTopicName()).thenReturn("real-time-topic"); + + // Test + String result = Utils.getRealTimeTopicNameFromStoreConfig(store); + + // Validate + assertEquals(result, "real-time-topic"); + + // Verify calls + verify(store).getHybridStoreConfig(); + verify(store).getName(); + verify(hybridStoreConfig).getRealTimeTopicName(); + } + + @Test + void testRealTimeTopicNameEmptyWithHybridConfig() { + // Mock the Store and HybridStoreConfig + Store store = mock(Store.class); + HybridStoreConfig hybridStoreConfig = mock(HybridStoreConfig.class); + + // Define behavior + when(store.getHybridStoreConfig()).thenReturn(hybridStoreConfig); + when(store.getName()).thenReturn("test-store"); + when(hybridStoreConfig.getRealTimeTopicName()).thenReturn(""); + + String result = Utils.getRealTimeTopicNameFromStoreConfig(store); + + assertEquals(result, "test-store_rt"); + } + + @Test + void testRealTimeTopicNameWithoutHybridConfig() { + // Mock the Store + Store store = mock(Store.class); + + // Define behavior + when(store.getHybridStoreConfig()).thenReturn(null); + when(store.getName()).thenReturn("test-store"); + + String result = Utils.getRealTimeTopicNameFromStoreConfig(store); + + assertEquals(result, "test-store_rt"); } @Test @@ -393,4 +450,87 @@ public void testGetLeaderTopicFromPubSubTopic() { Utils.resolveLeaderTopicFromPubSubTopic(pubSubTopicRepository, separateRealTimeTopic), realTimeTopic); } + + @Test + void testValidOldNameWithVersionIncrement() { + String oldName = "storeName_v1_rt"; + String expectedNewName = "storeName_v2_rt"; + + String result = Utils.createNewRealTimeTopicName(oldName); + + assertEquals(result, expectedNewName); + } + + @Test + void testWithVersionIncrement() { + String oldName = "storeName_v11_v55_rt"; + String expectedNewName = "storeName_v11_v56_rt"; + + String result = Utils.createNewRealTimeTopicName(oldName); + + assertEquals(result, expectedNewName); + } + + @Test + void testValidOldNameStartingNewVersion() { + String oldName = "storeName_rt"; + String expectedNewName = "storeName_v2_rt"; + + String result = Utils.createNewRealTimeTopicName(oldName); + + assertEquals(result, expectedNewName); + } + + @Test + void testInvalidOldNameNull() { + assertThrows(IllegalArgumentException.class, () -> Utils.createNewRealTimeTopicName(null)); + } + + @Test + void testInvalidOldNameWithoutSuffix() { + String oldName = "storeName_v1"; + assertThrows(IllegalArgumentException.class, () -> Utils.createNewRealTimeTopicName(oldName)); + } + + @Test + void testInvalidOldNameIncorrectFormat() { + String oldName = "storeName_v1_rt_extra"; + assertThrows(IllegalArgumentException.class, () -> Utils.createNewRealTimeTopicName(oldName)); + } + + @Test + void testInvalidOldNameWithNonNumericVersion() { + String oldName = "storeName_vX_rt"; + assertThrows(NumberFormatException.class, () -> Utils.createNewRealTimeTopicName(oldName)); + } + + @DataProvider(name = "booleanParsingData") + public Object[][] booleanParsingData() { + return new Object[][] { + // Valid cases + { "true", "testField", true }, // Valid "true" + { "false", "testField", false }, // Valid "false" + { "TRUE", "testField", true }, // Valid case-insensitive "TRUE" + { "FALSE", "testField", false }, // Valid case-insensitive "FALSE" + + // Invalid cases + { "notABoolean", "testField", null }, // Invalid string + { "123", "testField", null }, // Non-boolean numeric string + { "", "testField", null }, // Empty string + { null, "testField", null }, // Null input + }; + } + + @Test(dataProvider = "booleanParsingData") + public void testParseBooleanFromString(String value, String fieldName, Boolean expectedResult) { + if (expectedResult != null) { + // For valid cases + boolean result = Utils.parseBooleanFromString(value, fieldName); + assertEquals((boolean) expectedResult, result, "Parsed boolean value does not match expected value."); + return; + } + VeniceHttpException e = + expectThrows(VeniceHttpException.class, () -> Utils.parseBooleanFromString(value, fieldName)); + assertEquals(e.getHttpStatusCode(), HttpStatus.SC_BAD_REQUEST, "Invalid status code."); + } } diff --git a/internal/venice-test-common/build.gradle b/internal/venice-test-common/build.gradle index 82b56f843f6..d25b3ee8e9f 100644 --- a/internal/venice-test-common/build.gradle +++ b/internal/venice-test-common/build.gradle @@ -75,6 +75,7 @@ dependencies { implementation project(':internal:alpini:netty4:alpini-netty4-base') implementation project(':internal:alpini:router:alpini-router-api') implementation project(':internal:alpini:router:alpini-router-base') + implementation project(':integrations:venice-duckdb') implementation('org.apache.helix:helix-core:1.4.1:jdk8') { exclude group: 'org.apache.helix' @@ -137,7 +138,10 @@ def integrationTestBuckets = [ "1000": [ "com.linkedin.davinci.*", "com.linkedin.venice.endToEnd.DaVinciClientDiskFullTest", - "com.linkedin.venice.endToEnd.DaVinciClientMemoryLimitTest"], + "com.linkedin.venice.endToEnd.DaVinciClientMemoryLimitTest", + "com.linkedin.venice.endToEnd.DaVinciClientRecordTransformerTest"], + "1001": [ + "com.linkedin.venice.endToEnd.DuckDBDaVinciRecordTransformerIntegrationTest"], "1010": [ "com.linkedin.venice.endToEnd.DaVinciClientTest"], "1020": [ @@ -178,8 +182,7 @@ def integrationTestBuckets = [ "com.linkedin.venice.endToEnd.TestVson*", "com.linkedin.venice.endToEnd.Push*"], "1210": [ - "com.linkedin.venice.hadoop.*", - "com.linkedin.venice.endToEnd.DaVinciClientRecordTransformerTest"], + "com.linkedin.venice.hadoop.*"], "1220": [ "com.linkedin.venice.endToEnd.TestPushJob*"], "1230": [ diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/consumer/TestChangelogConsumer.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/consumer/TestChangelogConsumer.java index d12b085e990..f0062c1e9b1 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/consumer/TestChangelogConsumer.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/consumer/TestChangelogConsumer.java @@ -112,7 +112,6 @@ public class TestChangelogConsumer { private String clusterName; private VeniceClusterWrapper clusterWrapper; private ControllerClient parentControllerClient; - private static final List SCHEMA_HISTORY = Arrays.asList( NAME_RECORD_V1_SCHEMA, NAME_RECORD_V2_SCHEMA, diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/AbstractTestVeniceHelixAdmin.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/AbstractTestVeniceHelixAdmin.java index 6a73a694466..4c98479879a 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/AbstractTestVeniceHelixAdmin.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/AbstractTestVeniceHelixAdmin.java @@ -16,6 +16,9 @@ import static com.linkedin.venice.ConfigKeys.TOPIC_CLEANUP_SLEEP_INTERVAL_BETWEEN_TOPIC_LIST_FETCH_MS; import static com.linkedin.venice.ConfigKeys.UNREGISTER_METRIC_FOR_DELETED_STORE_ENABLED; import static com.linkedin.venice.ConfigKeys.ZOOKEEPER_ADDRESS; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; import com.linkedin.venice.common.VeniceSystemStoreUtils; import com.linkedin.venice.controller.kafka.TopicCleanupService; @@ -30,6 +33,8 @@ import com.linkedin.venice.integration.utils.ZkServerWrapper; import com.linkedin.venice.meta.Store; import com.linkedin.venice.pubsub.PubSubTopicRepository; +import com.linkedin.venice.pubsub.api.PubSubTopic; +import com.linkedin.venice.pubsub.manager.TopicManager; import com.linkedin.venice.stats.HelixMessageChannelStats; import com.linkedin.venice.utils.HelixUtils; import com.linkedin.venice.utils.MockTestStateModelFactory; @@ -90,7 +95,9 @@ public void setupCluster(boolean createParticipantStore, MetricsRepository metri pubSubBrokerWrapper = ServiceFactory.getPubSubBroker(); clusterName = Utils.getUniqueString("test-cluster"); Properties properties = getControllerProperties(clusterName); - if (!createParticipantStore) { + if (createParticipantStore) { + properties.put(PARTICIPANT_MESSAGE_STORE_ENABLED, true); + } else { properties.put(PARTICIPANT_MESSAGE_STORE_ENABLED, false); properties.put(ADMIN_HELIX_MESSAGING_CHANNEL_ENABLED, true); } @@ -259,17 +266,16 @@ VeniceHelixAdmin getFollower(List admins, String cluster) { * Participant store should be set up by child controller. */ private void verifyParticipantMessageStoreSetup() { + TopicManager topicManager = veniceAdmin.getTopicManager(); String participantStoreName = VeniceSystemStoreUtils.getParticipantStoreNameForCluster(clusterName); + PubSubTopic participantStoreRt = pubSubTopicRepository.getTopic(Utils.composeRealTimeTopic(participantStoreName)); TestUtils.waitForNonDeterministicAssertion(5, TimeUnit.SECONDS, () -> { Store store = veniceAdmin.getStore(clusterName, participantStoreName); - Assert.assertNotNull(store); - Assert.assertEquals(store.getVersions().size(), 1); + assertNotNull(store); + assertEquals(store.getVersions().size(), 1); + }); + TestUtils.waitForNonDeterministicAssertion(60, TimeUnit.SECONDS, () -> { + assertTrue(topicManager.containsTopic(participantStoreRt)); }); - TestUtils.waitForNonDeterministicAssertion( - 3, - TimeUnit.SECONDS, - () -> Assert.assertEquals( - veniceAdmin.getRealTimeTopic(clusterName, participantStoreName), - Utils.composeRealTimeTopic(participantStoreName))); } } diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestControllerSecureGrpcServer.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestControllerSecureGrpcServer.java new file mode 100644 index 00000000000..99f18d21300 --- /dev/null +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestControllerSecureGrpcServer.java @@ -0,0 +1,81 @@ +package com.linkedin.venice.controller; + +import static com.linkedin.venice.controller.server.grpc.ControllerGrpcSslSessionInterceptor.CLIENT_CERTIFICATE_CONTEXT_KEY; +import static org.testng.Assert.assertEquals; + +import com.linkedin.venice.controller.server.grpc.ControllerGrpcSslSessionInterceptor; +import com.linkedin.venice.grpc.GrpcUtils; +import com.linkedin.venice.grpc.VeniceGrpcServer; +import com.linkedin.venice.grpc.VeniceGrpcServerConfig; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcRequest; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcResponse; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcServiceGrpc; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcServiceGrpc.VeniceControllerGrpcServiceBlockingStub; +import com.linkedin.venice.security.SSLFactory; +import com.linkedin.venice.utils.SslUtils; +import com.linkedin.venice.utils.TestUtils; +import io.grpc.ChannelCredentials; +import io.grpc.Context; +import io.grpc.Grpc; +import io.grpc.ManagedChannel; +import java.security.cert.X509Certificate; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + + +public class TestControllerSecureGrpcServer { + private VeniceGrpcServer grpcSecureServer; + private int grpcSecureServerPort; + private SSLFactory sslFactory; + + @BeforeClass(alwaysRun = true) + public void setUpClass() { + sslFactory = SslUtils.getVeniceLocalSslFactory(); + grpcSecureServerPort = TestUtils.getFreePort(); + VeniceGrpcServerConfig grpcSecureServerConfig = + new VeniceGrpcServerConfig.Builder().setService(new VeniceControllerGrpcSecureServiceTestImpl()) + .setPort(grpcSecureServerPort) + .setNumThreads(2) + .setSslFactory(sslFactory) + .setInterceptor(new ControllerGrpcSslSessionInterceptor()) + .build(); + grpcSecureServer = new VeniceGrpcServer(grpcSecureServerConfig); + grpcSecureServer.start(); + } + + @AfterClass + public void tearDownClass() { + if (grpcSecureServer != null) { + grpcSecureServer.stop(); + } + } + + public static class VeniceControllerGrpcSecureServiceTestImpl + extends VeniceControllerGrpcServiceGrpc.VeniceControllerGrpcServiceImplBase { + @Override + public void discoverClusterForStore( + DiscoverClusterGrpcRequest request, + io.grpc.stub.StreamObserver responseObserver) { + X509Certificate clientCert = CLIENT_CERTIFICATE_CONTEXT_KEY.get(Context.current()); + if (clientCert == null) { + throw new RuntimeException("Client cert is null"); + } + DiscoverClusterGrpcResponse discoverClusterGrpcResponse = + DiscoverClusterGrpcResponse.newBuilder().setClusterName("test-cluster").build(); + responseObserver.onNext(discoverClusterGrpcResponse); + responseObserver.onCompleted(); + } + } + + @Test + public void testSslCertificatePropagationByGrpcInterceptor() { + String serverAddress = String.format("localhost:%d", grpcSecureServerPort); + ChannelCredentials credentials = GrpcUtils.buildChannelCredentials(sslFactory); + ManagedChannel channel = Grpc.newChannelBuilder(serverAddress, credentials).build(); + VeniceControllerGrpcServiceBlockingStub blockingStub = VeniceControllerGrpcServiceGrpc.newBlockingStub(channel); + DiscoverClusterGrpcRequest request = DiscoverClusterGrpcRequest.newBuilder().setStoreName("test-store").build(); + DiscoverClusterGrpcResponse response = blockingStub.discoverClusterForStore(request); + assertEquals(response.getClusterName(), "test-cluster"); + } +} diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestDeleteStoreDeletesRealtimeTopic.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestDeleteStoreDeletesRealtimeTopic.java index 0ad26693291..d25262c69d4 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestDeleteStoreDeletesRealtimeTopic.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestDeleteStoreDeletesRealtimeTopic.java @@ -15,6 +15,7 @@ import com.linkedin.venice.controllerapi.UpdateStoreQueryParams; import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.integration.utils.ServiceFactory; +import com.linkedin.venice.integration.utils.VeniceClusterCreateOptions; import com.linkedin.venice.integration.utils.VeniceClusterWrapper; import com.linkedin.venice.meta.StoreInfo; import com.linkedin.venice.meta.Version; @@ -40,48 +41,46 @@ public class TestDeleteStoreDeletesRealtimeTopic { private static final Logger LOGGER = LogManager.getLogger(TestDeleteStoreDeletesRealtimeTopic.class); - private VeniceClusterWrapper venice = null; - private AvroGenericStoreClient client = null; + private VeniceClusterWrapper veniceCluster = null; private ControllerClient controllerClient = null; private TopicManagerRepository topicManagerRepository = null; - private String storeName = null; private final PubSubTopicRepository pubSubTopicRepository = new PubSubTopicRepository(); - @BeforeClass + @BeforeClass(alwaysRun = true) public void setUp() { - venice = ServiceFactory.getVeniceCluster(); - controllerClient = - ControllerClient.constructClusterControllerClient(venice.getClusterName(), venice.getRandomRouterURL()); + veniceCluster = ServiceFactory.getVeniceCluster( + new VeniceClusterCreateOptions.Builder().numberOfControllers(1).numberOfServers(1).numberOfRouters(1).build()); + controllerClient = ControllerClient + .constructClusterControllerClient(veniceCluster.getClusterName(), veniceCluster.getRandomRouterURL()); topicManagerRepository = IntegrationTestPushUtils.getTopicManagerRepo( PUBSUB_OPERATION_TIMEOUT_MS_DEFAULT_VALUE, 100, 0l, - venice.getPubSubBrokerWrapper(), + veniceCluster.getPubSubBrokerWrapper(), pubSubTopicRepository); - storeName = Utils.getUniqueString("hybrid-store"); - venice.getNewStore(storeName); - makeStoreHybrid(venice, storeName, 100L, 5L); - client = ClientFactory.getAndStartGenericAvroClient( - ClientConfig.defaultGenericClientConfig(storeName).setVeniceURL(venice.getRandomRouterURL())); } - @AfterClass + @AfterClass(alwaysRun = true) public void cleanUp() { Utils.closeQuietlyWithErrorLogged(topicManagerRepository); - Utils.closeQuietlyWithErrorLogged(client); - Utils.closeQuietlyWithErrorLogged(venice); + Utils.closeQuietlyWithErrorLogged(veniceCluster); Utils.closeQuietlyWithErrorLogged(controllerClient); } @Test(timeOut = 60 * Time.MS_PER_SECOND) public void deletingHybridStoreDeletesRealtimeTopic() { - TestUtils.assertCommand(controllerClient.emptyPush(storeName, Utils.getUniqueString("push-id"), 1L)); + String storeName = Utils.getUniqueString("hybrid-store"); + veniceCluster.getNewStore(storeName); + makeStoreHybrid(veniceCluster, storeName, 100L, 5L); + + TestUtils + .assertCommand(controllerClient.sendEmptyPushAndWait(storeName, Utils.getUniqueString("push-id"), 1000, 60000)); // write streaming records SystemProducer veniceProducer = null; try { - veniceProducer = getSamzaProducer(venice, storeName, Version.PushType.STREAM); + veniceProducer = getSamzaProducer(veniceCluster, storeName, Version.PushType.STREAM); for (int i = 1; i <= 10; i++) { sendStreamingRecord(veniceProducer, storeName, i); } @@ -98,13 +97,16 @@ public void deletingHybridStoreDeletesRealtimeTopic() { assertEquals(storeResponse.getStore().getCurrentVersion(), 1, "The empty push has not activated yet..."); }); - TestUtils.waitForNonDeterministicAssertion(10, TimeUnit.SECONDS, () -> { - try { - assertEquals(client.get("9").get(), new Utf8("stream_9")); - } catch (Exception e) { - throw new VeniceException(e); - } - }); + try (AvroGenericStoreClient client = ClientFactory.getAndStartGenericAvroClient( + ClientConfig.defaultGenericClientConfig(storeName).setVeniceURL(veniceCluster.getRandomRouterURL()))) { + TestUtils.waitForNonDeterministicAssertion(10, TimeUnit.SECONDS, () -> { + try { + assertEquals(client.get("9").get(), new Utf8("stream_9")); + } catch (Exception e) { + throw new VeniceException(e); + } + }); + } // verify realtime topic exists PubSubTopic rtTopic = pubSubTopicRepository.getTopic(Utils.getRealTimeTopicName(storeInfo.get())); diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestHybridStoreRepartitioningWithMultiDataCenter.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestHybridStoreRepartitioningWithMultiDataCenter.java new file mode 100644 index 00000000000..2b8088cdb35 --- /dev/null +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestHybridStoreRepartitioningWithMultiDataCenter.java @@ -0,0 +1,157 @@ +package com.linkedin.venice.controller; + +import static com.linkedin.venice.ConfigKeys.CONTROLLER_ENABLE_HYBRID_STORE_PARTITION_COUNT_UPDATE; +import static com.linkedin.venice.ConfigKeys.DEFAULT_MAX_NUMBER_OF_PARTITIONS; +import static com.linkedin.venice.ConfigKeys.DEFAULT_NUMBER_OF_PARTITION_FOR_HYBRID; +import static com.linkedin.venice.ConfigKeys.DEFAULT_PARTITION_SIZE; + +import com.linkedin.venice.controllerapi.ControllerClient; +import com.linkedin.venice.controllerapi.NewStoreResponse; +import com.linkedin.venice.controllerapi.UpdateStoreQueryParams; +import com.linkedin.venice.integration.utils.ServiceFactory; +import com.linkedin.venice.integration.utils.VeniceMultiClusterWrapper; +import com.linkedin.venice.integration.utils.VeniceTwoLayerMultiRegionMultiClusterWrapper; +import com.linkedin.venice.meta.BackupStrategy; +import com.linkedin.venice.meta.StoreInfo; +import com.linkedin.venice.pubsub.PubSubTopicRepository; +import com.linkedin.venice.pubsub.api.PubSubTopic; +import com.linkedin.venice.pubsub.manager.TopicManager; +import com.linkedin.venice.utils.TestUtils; +import com.linkedin.venice.utils.TestWriteUtils; +import com.linkedin.venice.utils.Time; +import com.linkedin.venice.utils.Utils; +import java.util.ArrayList; +import java.util.List; +import java.util.Optional; +import java.util.Properties; +import java.util.concurrent.TimeUnit; +import java.util.stream.IntStream; +import org.testng.Assert; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + + +public class TestHybridStoreRepartitioningWithMultiDataCenter { + private static final int TEST_TIMEOUT = 90_000; // ms + private static final int NUMBER_OF_CHILD_DATACENTERS = 2; + private static final int NUMBER_OF_CLUSTERS = 1; + private static final String[] CLUSTER_NAMES = + IntStream.range(0, NUMBER_OF_CLUSTERS).mapToObj(i -> "venice-cluster" + i).toArray(String[]::new); + private List childDatacenters; + private VeniceTwoLayerMultiRegionMultiClusterWrapper multiRegionMultiClusterWrapper; + List topicManagers; + PubSubTopicRepository pubSubTopicRepository = new PubSubTopicRepository(); + + @BeforeClass + public void setUp() { + Properties controllerProps = new Properties(); + controllerProps.put(DEFAULT_NUMBER_OF_PARTITION_FOR_HYBRID, 2); + controllerProps.put(DEFAULT_MAX_NUMBER_OF_PARTITIONS, 3); + controllerProps.put(DEFAULT_PARTITION_SIZE, 1024); + controllerProps.put(CONTROLLER_ENABLE_HYBRID_STORE_PARTITION_COUNT_UPDATE, true); + multiRegionMultiClusterWrapper = ServiceFactory.getVeniceTwoLayerMultiRegionMultiClusterWrapper( + NUMBER_OF_CHILD_DATACENTERS, + NUMBER_OF_CLUSTERS, + 1, + 1, + 1, + 1, + 1, + Optional.of(controllerProps), + Optional.of(controllerProps), + Optional.empty()); + + childDatacenters = multiRegionMultiClusterWrapper.getChildRegions(); + topicManagers = new ArrayList<>(2); + topicManagers + .add(childDatacenters.get(0).getControllers().values().iterator().next().getVeniceAdmin().getTopicManager()); + topicManagers + .add(childDatacenters.get(1).getControllers().values().iterator().next().getVeniceAdmin().getTopicManager()); + } + + @AfterClass(alwaysRun = true) + public void cleanUp() { + Utils.closeQuietlyWithErrorLogged(multiRegionMultiClusterWrapper); + } + + @Test(timeOut = TEST_TIMEOUT) + public void testHybridStoreRepartitioning() { + String storeName = Utils.getUniqueString("TestHybridStoreRepartitioning"); + String clusterName = CLUSTER_NAMES[0]; + String parentControllerURLs = multiRegionMultiClusterWrapper.getControllerConnectString(); + ControllerClient parentControllerClient = + ControllerClient.constructClusterControllerClient(clusterName, parentControllerURLs); + ControllerClient[] childControllerClients = new ControllerClient[childDatacenters.size()]; + for (int i = 0; i < childDatacenters.size(); i++) { + childControllerClients[i] = + new ControllerClient(clusterName, childDatacenters.get(i).getControllerConnectString()); + } + + NewStoreResponse newStoreResponse = + parentControllerClient.retryableRequest(5, c -> c.createNewStore(storeName, "", "\"string\"", "\"string\"")); + Assert.assertFalse( + newStoreResponse.isError(), + "The NewStoreResponse returned an error: " + newStoreResponse.getError()); + + UpdateStoreQueryParams updateStoreParams = new UpdateStoreQueryParams(); + updateStoreParams.setIncrementalPushEnabled(true) + .setBackupStrategy(BackupStrategy.KEEP_MIN_VERSIONS) + .setNumVersionsToPreserve(2) + .setHybridRewindSeconds(1000) + .setActiveActiveReplicationEnabled(true) + .setHybridOffsetLagThreshold(1000); + TestWriteUtils.updateStore(storeName, parentControllerClient, updateStoreParams); + + // create new version by doing an empty push + parentControllerClient + .sendEmptyPushAndWait(storeName, Utils.getUniqueString("empty-push"), 1L, 60L * Time.MS_PER_SECOND); + + for (ControllerClient controllerClient: childControllerClients) { + Assert.assertEquals(controllerClient.getStore(storeName).getStore().getCurrentVersion(), 1); + } + + TestWriteUtils.updateStore(storeName, parentControllerClient, new UpdateStoreQueryParams().setPartitionCount(2)); + + for (int i = 0; i < childControllerClients.length; i++) { + final int index = i; + TestUtils.waitForNonDeterministicAssertion(30, TimeUnit.SECONDS, () -> { + StoreInfo storeInfo = childControllerClients[index].getStore(storeName).getStore(); + String realTimeTopicNameInVersion = Utils.getRealTimeTopicName(storeInfo.getVersions().get(0)); + PubSubTopic realTimePubSubTopic = pubSubTopicRepository.getTopic(realTimeTopicNameInVersion); + + // verify rt topic is created with the default partition count = 3 + Assert.assertEquals(topicManagers.get(index).getPartitionCount(realTimePubSubTopic), 3); + }); + } + + // create new version by doing an empty push + parentControllerClient + .sendEmptyPushAndWait(storeName, Utils.getUniqueString("empty-push"), 1L, 60L * Time.MS_PER_SECOND); + + for (ControllerClient controllerClient: childControllerClients) { + Assert.assertEquals(controllerClient.getStore(storeName).getStore().getCurrentVersion(), 2); + } + + for (int i = 0; i < childControllerClients.length; i++) { + final int idx = i; + TestUtils.waitForNonDeterministicAssertion(30, TimeUnit.SECONDS, () -> { + StoreInfo storeInfo = childControllerClients[idx].getStore(storeName).getStore(); + String realTimeTopicNameInBackupVersion = Utils.getRealTimeTopicName(storeInfo.getVersions().get(0)); + String realTimeTopicNameInCurrentVersion = Utils.getRealTimeTopicName(storeInfo.getVersions().get(1)); + String expectedRealTimeTopicNameInStoreConfig = + Utils.createNewRealTimeTopicName(realTimeTopicNameInBackupVersion); + String actualRealTimeTopicNameInStoreConfig = storeInfo.getHybridStoreConfig().getRealTimeTopicName(); + PubSubTopic newRtPubSubTopic = pubSubTopicRepository.getTopic(realTimeTopicNameInCurrentVersion); + + // verify rt topic name + Assert.assertNotEquals(realTimeTopicNameInBackupVersion, realTimeTopicNameInCurrentVersion); + Assert.assertEquals(realTimeTopicNameInCurrentVersion, actualRealTimeTopicNameInStoreConfig); + Assert.assertEquals(actualRealTimeTopicNameInStoreConfig, expectedRealTimeTopicNameInStoreConfig); + + // verify rt topic is created with the updated partition count = 2 + Assert.assertEquals(topicManagers.get(idx).getPartitionCount(newRtPubSubTopic), 2); + }); + } + } +} diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestInstanceRemovable.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestInstanceRemovable.java index a6121f6e132..17e5ddb70d0 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestInstanceRemovable.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestInstanceRemovable.java @@ -1,10 +1,12 @@ package com.linkedin.venice.controller; +import com.linkedin.venice.ConfigKeys; import com.linkedin.venice.controllerapi.ControllerClient; import com.linkedin.venice.controllerapi.StoppableNodeStatusResponse; import com.linkedin.venice.controllerapi.UpdateStoreQueryParams; import com.linkedin.venice.controllerapi.VersionCreationResponse; import com.linkedin.venice.integration.utils.ServiceFactory; +import com.linkedin.venice.integration.utils.VeniceClusterCreateOptions; import com.linkedin.venice.integration.utils.VeniceClusterWrapper; import com.linkedin.venice.integration.utils.VeniceServerWrapper; import com.linkedin.venice.meta.PartitionAssignment; @@ -34,17 +36,16 @@ public class TestInstanceRemovable { int replicaFactor = 3; private void setupCluster(int numberOfServer) { - int numberOfController = 1; - int numberOfRouter = 1; - + Properties properties = new Properties(); + properties.setProperty(ConfigKeys.PARTICIPANT_MESSAGE_STORE_ENABLED, "false"); cluster = ServiceFactory.getVeniceCluster( - numberOfController, - numberOfServer, - numberOfRouter, - replicaFactor, - partitionSize, - false, - false); + new VeniceClusterCreateOptions.Builder().numberOfControllers(1) + .numberOfServers(numberOfServer) + .numberOfRouters(1) + .replicationFactor(replicaFactor) + .partitionSize(partitionSize) + .extraProperties(properties) + .build()); } @AfterMethod @@ -169,7 +170,7 @@ public void testIsInstanceRemovableAfterPush() throws Exception { // Wait push completed. TestUtils.waitForNonDeterministicCompletion( - 3, + 30, TimeUnit.SECONDS, () -> cluster.getLeaderVeniceController() .getVeniceAdmin() diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestVeniceHelixAdminWithIsolatedEnvironment.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestVeniceHelixAdminWithIsolatedEnvironment.java index 63cbfe22e50..5321196dbec 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestVeniceHelixAdminWithIsolatedEnvironment.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestVeniceHelixAdminWithIsolatedEnvironment.java @@ -131,7 +131,7 @@ public void testGetLeaderController() { Assert.assertEquals( veniceAdmin.getLeaderController(clusterName).getNodeId(), Utils.getHelixNodeIdentifier(controllerConfig.getAdminHostname(), controllerConfig.getAdminPort())); - // Create a new controller and test getLeaderController again. + // Create a new controller and test getLeaderControllerDetails again. int newAdminPort = controllerConfig.getAdminPort() - 10; PropertyBuilder builder = new PropertyBuilder().put(controllerProps.toProperties()).put("admin.port", newAdminPort); VeniceProperties newControllerProps = builder.build(); diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestVeniceHelixAdminWithSharedEnvironment.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestVeniceHelixAdminWithSharedEnvironment.java index cfb1e8cb9e6..0fffecfb67f 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestVeniceHelixAdminWithSharedEnvironment.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/TestVeniceHelixAdminWithSharedEnvironment.java @@ -11,9 +11,12 @@ import static org.mockito.Mockito.eq; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.reset; +import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; import com.linkedin.venice.ConfigKeys; +import com.linkedin.venice.common.VeniceSystemStoreType; import com.linkedin.venice.common.VeniceSystemStoreUtils; import com.linkedin.venice.compression.CompressionStrategy; import com.linkedin.venice.controller.exception.HelixClusterMaintenanceModeException; @@ -71,6 +74,7 @@ import java.util.HashMap; import java.util.List; import java.util.Map; +import java.util.Objects; import java.util.Optional; import java.util.Properties; import java.util.concurrent.CompletableFuture; @@ -641,34 +645,46 @@ public void testAddVersionWhenClusterInMaintenanceMode() { } @Test(timeOut = TOTAL_TIMEOUT_FOR_LONG_TEST_MS) - public void testGetRealTimeTopic() { + public void testEnsureRealTimeTopicExistsForUserSystemStores() { String storeName = Utils.getUniqueString("store"); + String metaStoreName = VeniceSystemStoreType.META_STORE.getSystemStoreName(storeName); - // Must not be able to get a real time topic until the store is created - Assert.assertThrows(VeniceNoStoreException.class, () -> veniceAdmin.getRealTimeTopic(clusterName, storeName)); + Exception notSystemStoreException = Assert.expectThrows( + VeniceNoStoreException.class, + () -> veniceAdmin.ensureRealTimeTopicExistsForUserSystemStores(clusterName, metaStoreName)); + assertTrue( + notSystemStoreException.getMessage().contains("does not exist in"), + "Got unexpected error message: " + notSystemStoreException.getMessage()); veniceAdmin.createStore(clusterName, storeName, "owner", KEY_SCHEMA, VALUE_SCHEMA); - Store store = veniceAdmin.getStore(clusterName, storeName); + Store userStore = veniceAdmin.getStore(clusterName, storeName); + assertNotNull(userStore, "User store should be created and not null"); veniceAdmin.updateStore( clusterName, storeName, - new UpdateStoreQueryParams().setHybridRewindSeconds(25L).setHybridOffsetLagThreshold(100L)); // make store - // hybrid + new UpdateStoreQueryParams().setHybridRewindSeconds(25L).setHybridOffsetLagThreshold(100L)); - try { - veniceAdmin.getRealTimeTopic(clusterName, storeName); - Assert.fail("Must not be able to get a real time topic until the store is initialized with a version"); - } catch (VeniceException e) { - Assert.assertTrue( - e.getMessage().contains("is not initialized with a version"), - "Got unexpected error message: " + e.getMessage()); - } - - int partitions = 2; // TODO verify partition count for RT topic. - veniceAdmin.incrementVersionIdempotent(clusterName, storeName, Version.guidBasedDummyPushId(), partitions, 1); + Exception exception = Assert.expectThrows( + VeniceException.class, + () -> veniceAdmin.ensureRealTimeTopicExistsForUserSystemStores(clusterName, storeName)); + assertTrue( + exception.getMessage().contains("not a user system store"), + "Got unexpected error message: " + notSystemStoreException.getMessage()); + + String pushStatusStoreName = VeniceSystemStoreType.DAVINCI_PUSH_STATUS_STORE.getSystemStoreName(storeName); + Store pushStatusStore = veniceAdmin.getStore(clusterName, pushStatusStoreName); + PubSubTopic pushStatusRealTimeTopic = pubSubTopicRepository.getTopic(Utils.getRealTimeTopicName(pushStatusStore)); + assertNotNull(pushStatusStore, "Push status store should not be created yet"); + TestUtils.waitForNonDeterministicCompletion( + 30, + TimeUnit.SECONDS, + () -> !veniceAdmin.getTopicManager().containsTopic(pushStatusRealTimeTopic)); - String rtTopic = veniceAdmin.getRealTimeTopic(clusterName, storeName); - Assert.assertEquals(rtTopic, Utils.getRealTimeTopicName(store)); + veniceAdmin.ensureRealTimeTopicExistsForUserSystemStores(clusterName, pushStatusStoreName); + TestUtils.waitForNonDeterministicCompletion( + 30, + TimeUnit.SECONDS, + () -> veniceAdmin.getTopicManager().containsTopic(pushStatusRealTimeTopic)); } @Test(timeOut = TOTAL_TIMEOUT_FOR_LONG_TEST_MS) @@ -1679,29 +1695,33 @@ public void testGetIncrementalPushVersion() { new UpdateStoreQueryParams().setHybridOffsetLagThreshold(1) .setHybridRewindSeconds(0) .setIncrementalPushEnabled(true)); - veniceAdmin.incrementVersionIdempotent( + Version version = veniceAdmin.incrementVersionIdempotent( clusterName, incrementalAndHybridEnabledStoreName, Version.guidBasedDummyPushId(), 1, 1); - String rtTopic = veniceAdmin.getRealTimeTopic(clusterName, incrementalAndHybridEnabledStoreName); TestUtils.waitForNonDeterministicCompletion( TOTAL_TIMEOUT_FOR_SHORT_TEST_MS, TimeUnit.MILLISECONDS, () -> veniceAdmin.getCurrentVersion(clusterName, incrementalAndHybridEnabledStoreName) == 1); + PubSubTopic rtTopic = pubSubTopicRepository.getTopic(Utils.getRealTimeTopicName(version)); + TestUtils.waitForNonDeterministicCompletion( + TOTAL_TIMEOUT_FOR_LONG_TEST_MS, + TimeUnit.MILLISECONDS, + () -> veniceAdmin.getTopicManager().containsTopic(rtTopic)); // For incremental push policy INCREMENTAL_PUSH_SAME_AS_REAL_TIME, incremental push should succeed even if version // topic is truncated veniceAdmin.truncateKafkaTopic(Version.composeKafkaTopic(incrementalAndHybridEnabledStoreName, 1)); - veniceAdmin.getIncrementalPushVersion(clusterName, incrementalAndHybridEnabledStoreName); + veniceAdmin.getIncrementalPushVersion(clusterName, incrementalAndHybridEnabledStoreName, "test-job-1"); // For incremental push policy INCREMENTAL_PUSH_SAME_AS_REAL_TIME, incremental push should fail if rt topic is // truncated - veniceAdmin.truncateKafkaTopic(rtTopic); + veniceAdmin.truncateKafkaTopic(rtTopic.getName()); Assert.assertThrows( VeniceException.class, - () -> veniceAdmin.getIncrementalPushVersion(clusterName, incrementalAndHybridEnabledStoreName)); + () -> veniceAdmin.getIncrementalPushVersion(clusterName, incrementalAndHybridEnabledStoreName, "test-job-1")); } @Test(timeOut = TOTAL_TIMEOUT_FOR_LONG_TEST_MS) @@ -1863,8 +1883,21 @@ public void testHybridStoreToBatchOnly() { Assert.assertTrue(veniceAdmin.getStore(clusterName, storeName).isSeparateRealTimeTopicEnabled()); Assert.assertTrue(veniceAdmin.getStore(clusterName, storeName).getVersion(1).isSeparateRealTimeTopicEnabled()); - String rtTopic = veniceAdmin.getRealTimeTopic(clusterName, storeName); - String incrementalPushRealTimeTopic = veniceAdmin.getSeparateRealTimeTopic(clusterName, storeName); + Store store = Objects.requireNonNull(veniceAdmin.getStore(clusterName, storeName), "Store should not be null"); + + String rtTopic = Utils.getRealTimeTopicName(store); + PubSubTopic rtPubSubTopic = pubSubTopicRepository.getTopic(rtTopic); + String incrementalPushRealTimeTopic = Version.composeSeparateRealTimeTopic(storeName); + PubSubTopic incrementalPushRealTimePubSubTopic = pubSubTopicRepository.getTopic(incrementalPushRealTimeTopic); + TestUtils.waitForNonDeterministicCompletion( + TOTAL_TIMEOUT_FOR_SHORT_TEST_MS, + TimeUnit.MILLISECONDS, + () -> veniceAdmin.getTopicManager().containsTopic(rtPubSubTopic)); + TestUtils.waitForNonDeterministicCompletion( + TOTAL_TIMEOUT_FOR_SHORT_TEST_MS, + TimeUnit.MILLISECONDS, + () -> veniceAdmin.getTopicManager().containsTopic(incrementalPushRealTimePubSubTopic)); + Assert.assertFalse(veniceAdmin.isTopicTruncated(rtTopic)); Assert.assertFalse(veniceAdmin.isTopicTruncated(incrementalPushRealTimeTopic)); veniceAdmin.updateStore( diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/VeniceParentHelixAdminTest.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/VeniceParentHelixAdminTest.java index cfcfb0d59e2..b1303d94c23 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/VeniceParentHelixAdminTest.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/VeniceParentHelixAdminTest.java @@ -10,6 +10,7 @@ import static com.linkedin.venice.controller.SchemaConstants.VALUE_SCHEMA_FOR_WRITE_COMPUTE_V3; import static com.linkedin.venice.controller.SchemaConstants.VALUE_SCHEMA_FOR_WRITE_COMPUTE_V4; import static com.linkedin.venice.controller.SchemaConstants.VALUE_SCHEMA_FOR_WRITE_COMPUTE_V5; +import static com.linkedin.venice.pubsub.PubSubConstants.PUBSUB_OPERATION_TIMEOUT_MS_DEFAULT_VALUE; import static com.linkedin.venice.utils.ByteUtils.BYTES_PER_MB; import static com.linkedin.venice.utils.TestUtils.assertCommand; import static com.linkedin.venice.utils.TestUtils.waitForNonDeterministicAssertion; @@ -31,6 +32,7 @@ import com.linkedin.venice.controllerapi.StoreResponse; import com.linkedin.venice.controllerapi.UpdateStoreQueryParams; import com.linkedin.venice.controllerapi.VersionCreationResponse; +import com.linkedin.venice.integration.utils.PubSubBrokerWrapper; import com.linkedin.venice.integration.utils.ServiceFactory; import com.linkedin.venice.integration.utils.VeniceClusterWrapper; import com.linkedin.venice.integration.utils.VeniceControllerWrapper; @@ -40,9 +42,14 @@ import com.linkedin.venice.meta.HybridStoreConfig; import com.linkedin.venice.meta.StoreInfo; import com.linkedin.venice.meta.Version; +import com.linkedin.venice.pubsub.PubSubTopicRepository; +import com.linkedin.venice.pubsub.api.PubSubTopic; +import com.linkedin.venice.pubsub.manager.TopicManager; +import com.linkedin.venice.pubsub.manager.TopicManagerRepository; import com.linkedin.venice.schema.AvroSchemaParseUtils; import com.linkedin.venice.schema.writecompute.WriteComputeSchemaConverter; import com.linkedin.venice.security.SSLFactory; +import com.linkedin.venice.utils.IntegrationTestPushUtils; import com.linkedin.venice.utils.SslUtils; import com.linkedin.venice.utils.TestUtils; import com.linkedin.venice.utils.TestWriteUtils; @@ -298,7 +305,24 @@ public void testResourceCleanupCheckForStoreRecreation() { TimeUnit.SECONDS); // Delete the store and try re-creation. - assertFalse(parentControllerClient.disableAndDeleteStore(storeName).isError(), "Delete store shouldn't fail"); + TestUtils.assertCommand(parentControllerClient.disableAndDeleteStore(storeName), "Delete store shouldn't fail"); + + PubSubBrokerWrapper parentPubSub = twoLayerMultiRegionMultiClusterWrapper.getParentKafkaBrokerWrapper(); + PubSubTopicRepository pubSubTopicRepository = new PubSubTopicRepository(); + // Manually create an RT topic in the parent region to simulate its presence for lingering system store resources. + // This is necessary because RT topics are no longer automatically created for regional system stores such as meta + // and ps3. + try (TopicManagerRepository topicManagerRepo = IntegrationTestPushUtils + .getTopicManagerRepo(PUBSUB_OPERATION_TIMEOUT_MS_DEFAULT_VALUE, 100, 0l, parentPubSub, pubSubTopicRepository); + TopicManager topicManager = topicManagerRepo.getLocalTopicManager()) { + PubSubTopic metaStoreRT = pubSubTopicRepository.getTopic(Version.composeRealTimeTopic(metaSystemStoreName)); + topicManager.createTopic(metaStoreRT, 1, 1, true); + TestUtils.waitForNonDeterministicAssertion( + 30, + TimeUnit.SECONDS, + () -> assertTrue(topicManager.containsTopic(metaStoreRT))); + } + // Re-create the same store right away will fail because of lingering system store resources controllerResponse = parentControllerClient.createNewStore(storeName, "test", "\"string\"", "\"string\""); assertTrue( diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/server/TestAdminSparkServerWithMultiServers.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/server/TestAdminSparkServerWithMultiServers.java index 0e8dc96bde9..1d52e718f60 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/server/TestAdminSparkServerWithMultiServers.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/server/TestAdminSparkServerWithMultiServers.java @@ -254,7 +254,9 @@ public void requestTopicIsIdempotent() { controllerClient.updateStore( storeName, new UpdateStoreQueryParams().setHybridRewindSeconds(1000).setHybridOffsetLagThreshold(1000)); - controllerClient.emptyPush(storeName, Utils.getUniqueString("emptyPushId"), 10000); + TestUtils.assertCommand( + controllerClient + .sendEmptyPushAndWait(storeName, Utils.getUniqueString("emptyPushId"), 10000, TEST_TIMEOUT)); } // Both diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/server/TestAdminSparkWithMocks.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/server/TestAdminSparkWithMocks.java index 925625ab4c5..81e539b60f6 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/server/TestAdminSparkWithMocks.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/controller/server/TestAdminSparkWithMocks.java @@ -1,14 +1,17 @@ package com.linkedin.venice.controller.server; import static com.linkedin.venice.meta.Store.NON_EXISTING_VERSION; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.any; import static org.mockito.Mockito.anyBoolean; import static org.mockito.Mockito.anyString; import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.eq; +import static org.mockito.Mockito.mock; +import com.linkedin.venice.controller.Admin; +import com.linkedin.venice.controller.ControllerRequestHandlerDependencies; import com.linkedin.venice.controller.ParentControllerRegionState; -import com.linkedin.venice.controller.VeniceHelixAdmin; import com.linkedin.venice.controllerapi.ControllerApiConstants; import com.linkedin.venice.controllerapi.ControllerRoute; import com.linkedin.venice.controllerapi.VersionCreationResponse; @@ -24,11 +27,11 @@ import com.linkedin.venice.meta.Store; import com.linkedin.venice.meta.Version; import com.linkedin.venice.meta.VersionImpl; +import com.linkedin.venice.meta.VersionStatus; import com.linkedin.venice.meta.ZKStore; import com.linkedin.venice.utils.DataProviderUtils; import com.linkedin.venice.utils.ObjectMapperFactory; import com.linkedin.venice.utils.SslUtils; -import com.linkedin.venice.utils.Utils; import java.util.ArrayList; import java.util.Arrays; import java.util.List; @@ -43,6 +46,7 @@ import org.apache.http.message.BasicNameValuePair; import org.mockito.Mockito; import org.testng.Assert; +import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -52,10 +56,20 @@ * verifying any state changes that would be triggered by the admin. */ public class TestAdminSparkWithMocks { + private VeniceControllerRequestHandler requestHandler; + private Admin admin; + + @BeforeMethod(alwaysRun = true) + public void setUp() { + admin = Mockito.mock(Admin.class); + ControllerRequestHandlerDependencies dependencies = mock(ControllerRequestHandlerDependencies.class); + doReturn(admin).when(dependencies).getAdmin(); + requestHandler = new VeniceControllerRequestHandler(dependencies); + } + @Test - public void testGetRealTimeTopicUsesAdmin() throws Exception { + public void testGetRealTimeTopicForStreamPushJobUsesAdmin() throws Exception { // setup server with mock admin, note returns topic "store_rt" - VeniceHelixAdmin admin = Mockito.mock(VeniceHelixAdmin.class); Store mockStore = new ZKStore( "store", "owner", @@ -72,17 +86,28 @@ public void testGetRealTimeTopicUsesAdmin() throws Exception { HybridStoreConfigImpl.DEFAULT_HYBRID_TIME_LAG_THRESHOLD, DataReplicationPolicy.NON_AGGREGATE, BufferReplayPolicy.REWIND_FROM_EOP)); + Version hybridVersion = new VersionImpl("store", 1, "pushJobId-1234", 33); + hybridVersion.setHybridStoreConfig(mockStore.getHybridStoreConfig()); + hybridVersion.setStatus(VersionStatus.ONLINE); + mockStore.addVersion(hybridVersion); + + // check store partition count is different from hybrid version partition count so that we can verify the + // partition count is updated to the hybrid version partition count in response + Assert.assertNotEquals(mockStore.getPartitionCount(), hybridVersion.getPartitionCount()); + doReturn(mockStore).when(admin).getStore(anyString(), anyString()); doReturn(true).when(admin).isLeaderControllerFor(anyString()); doReturn(1).when(admin).getReplicationFactor(anyString(), anyString()); doReturn(1).when(admin).calculateNumberOfPartitions(anyString(), anyString()); doReturn("kafka-bootstrap").when(admin).getKafkaBootstrapServers(anyBoolean()); - doReturn("store_rt").when(admin).getRealTimeTopic(anyString(), anyString()); - doReturn("store_rt").when(admin).getRealTimeTopic(anyString(), any(Store.class)); + doReturn(hybridVersion).when(admin).getReferenceVersionForStreamingWrites(anyString(), anyString(), any()); // Add a banned route not relevant to the test just to make sure theres coverage for unbanned routes still be // accessible - AdminSparkServer server = - ServiceFactory.getMockAdminSparkServer(admin, "clustername", Arrays.asList(ControllerRoute.ADD_DERIVED_SCHEMA)); + AdminSparkServer server = ServiceFactory.getMockAdminSparkServer( + admin, + "clustername", + Arrays.asList(ControllerRoute.ADD_DERIVED_SCHEMA), + requestHandler); int port = server.getPort(); // build request @@ -110,6 +135,7 @@ public void testGetRealTimeTopicUsesAdmin() throws Exception { // verify response, note we expect same topic, "store_rt" Assert.assertFalse(responseObject.isError(), "unexpected error: " + responseObject.getError()); Assert.assertEquals(responseObject.getKafkaTopic(), "store_rt"); + Assert.assertEquals(responseObject.getPartitions(), hybridVersion.getPartitionCount()); server.stop(); } @@ -117,7 +143,6 @@ public void testGetRealTimeTopicUsesAdmin() throws Exception { @Test public void testBannedRoutesAreRejected() throws Exception { // setup server with mock admin, note returns topic "store_rt" - VeniceHelixAdmin admin = Mockito.mock(VeniceHelixAdmin.class); Store mockStore = new ZKStore( "store", "owner", @@ -139,10 +164,8 @@ public void testBannedRoutesAreRejected() throws Exception { doReturn(1).when(admin).getReplicationFactor(anyString(), anyString()); doReturn(1).when(admin).calculateNumberOfPartitions(anyString(), anyString()); doReturn("kafka-bootstrap").when(admin).getKafkaBootstrapServers(anyBoolean()); - doReturn("store_rt").when(admin).getRealTimeTopic(anyString(), anyString()); - doReturn("store_rt").when(admin).getRealTimeTopic(anyString(), any(Store.class)); - AdminSparkServer server = - ServiceFactory.getMockAdminSparkServer(admin, "clustername", Arrays.asList(ControllerRoute.REQUEST_TOPIC)); + AdminSparkServer server = ServiceFactory + .getMockAdminSparkServer(admin, "clustername", Arrays.asList(ControllerRoute.REQUEST_TOPIC), requestHandler); int port = server.getPort(); // build request @@ -189,7 +212,6 @@ public void testAAIncrementalPushRTSourceRegion(boolean sourceGridFabricPresent, Optional optionalemergencySourceRegion = Optional.empty(); Optional optionalSourceGridSourceFabric = Optional.empty(); - VeniceHelixAdmin admin = Mockito.mock(VeniceHelixAdmin.class); Store mockStore = new ZKStore( storeName, "owner", @@ -217,8 +239,6 @@ public void testAAIncrementalPushRTSourceRegion(boolean sourceGridFabricPresent, doReturn(corpRegionKafka).when(admin).getKafkaBootstrapServers(anyBoolean()); doReturn(true).when(admin).whetherEnableBatchPushFromAdmin(anyString()); doReturn(true).when(admin).isActiveActiveReplicationEnabledInAllRegion(clusterName, storeName, false); - doReturn(Utils.getRealTimeTopicName(mockStore)).when(admin).getRealTimeTopic(anyString(), anyString()); - doReturn(Utils.getRealTimeTopicName(mockStore)).when(admin).getRealTimeTopic(anyString(), any(Store.class)); doReturn(corpRegionKafka).when(admin).getNativeReplicationKafkaBootstrapServerAddress(corpRegion); doReturn(emergencySourceRegionKafka).when(admin) .getNativeReplicationKafkaBootstrapServerAddress(emergencySourceRegion); @@ -265,8 +285,11 @@ public void testAAIncrementalPushRTSourceRegion(boolean sourceGridFabricPresent, // Add a banned route not relevant to the test just to make sure theres coverage for unbanned routes still be // accessible - AdminSparkServer server = - ServiceFactory.getMockAdminSparkServer(admin, "clustername", Arrays.asList(ControllerRoute.ADD_DERIVED_SCHEMA)); + AdminSparkServer server = ServiceFactory.getMockAdminSparkServer( + admin, + "clustername", + Arrays.asList(ControllerRoute.ADD_DERIVED_SCHEMA), + requestHandler); int port = server.getPort(); final HttpPost post = new HttpPost("http://localhost:" + port + ControllerRoute.REQUEST_TOPIC.getPath()); post.setEntity(new UrlEncodedFormEntity(params)); @@ -303,7 +326,6 @@ public void testAAIncrementalPushRTSourceRegion(boolean sourceGridFabricPresent, public void testSamzaReplicationPolicyMode(boolean samzaPolicy, boolean storePolicy, boolean aaEnabled) throws Exception { // setup server with mock admin, note returns topic "store_rt" - VeniceHelixAdmin admin = Mockito.mock(VeniceHelixAdmin.class); Store mockStore = new ZKStore( "store", "owner", @@ -330,22 +352,29 @@ public void testSamzaReplicationPolicyMode(boolean samzaPolicy, boolean storePol DataReplicationPolicy.NON_AGGREGATE, BufferReplayPolicy.REWIND_FROM_EOP)); } + Version hybridVersion = new VersionImpl("store", 1, "pushJobId-1234", 33); + hybridVersion.setHybridStoreConfig(mockStore.getHybridStoreConfig()); + hybridVersion.setStatus(VersionStatus.ONLINE); + mockStore.addVersion(hybridVersion); + doReturn(mockStore).when(admin).getStore(anyString(), anyString()); doReturn(true).when(admin).isLeaderControllerFor(anyString()); doReturn(1).when(admin).getReplicationFactor(anyString(), anyString()); doReturn(1).when(admin).calculateNumberOfPartitions(anyString(), anyString()); doReturn("kafka-bootstrap").when(admin).getKafkaBootstrapServers(anyBoolean()); - doReturn("store_rt").when(admin).getRealTimeTopic(anyString(), anyString()); - doReturn("store_rt").when(admin).getRealTimeTopic(anyString(), any(Store.class)); + doReturn(hybridVersion).when(admin).getReferenceVersionForStreamingWrites(anyString(), anyString(), anyString()); doReturn(samzaPolicy).when(admin).isParent(); doReturn(ParentControllerRegionState.ACTIVE).when(admin).getParentControllerRegionState(); doReturn(aaEnabled).when(admin).isActiveActiveReplicationEnabledInAllRegion(anyString(), anyString(), eq(true)); mockStore.setActiveActiveReplicationEnabled(aaEnabled); - // Add a banned route not relevant to the test just to make sure theres coverage for unbanned routes still be + // Add a banned route not relevant to the test just to make sure there is coverage for unbanned routes still be // accessible - AdminSparkServer server = - ServiceFactory.getMockAdminSparkServer(admin, "clustername", Arrays.asList(ControllerRoute.ADD_DERIVED_SCHEMA)); + AdminSparkServer server = ServiceFactory.getMockAdminSparkServer( + admin, + "clustername", + Arrays.asList(ControllerRoute.ADD_DERIVED_SCHEMA), + requestHandler); int port = server.getPort(); // build request diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/BlobP2PTransferAmongServersTest.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/BlobP2PTransferAmongServersTest.java index 621e4a8a6ee..a997459055e 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/BlobP2PTransferAmongServersTest.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/BlobP2PTransferAmongServersTest.java @@ -37,7 +37,6 @@ import org.apache.samza.system.SystemProducer; import org.testng.Assert; import org.testng.annotations.AfterMethod; -import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -47,13 +46,10 @@ public class BlobP2PTransferAmongServersTest { private static final int STREAMING_RECORD_SIZE = 1024; private String path1; private String path2; + private int server1Port; + private int server2Port; private VeniceClusterWrapper cluster; - @BeforeMethod(alwaysRun = true) - public void setUp() { - cluster = initializeVeniceCluster(); - } - @AfterMethod(alwaysRun = true) public void tearDown() { Utils.closeQuietlyWithErrorLogged(cluster); @@ -68,12 +64,14 @@ public void tearDown() { @Test(singleThreaded = true, timeOut = 180000) public void testBlobP2PTransferAmongServersForBatchStore() throws Exception { + cluster = initializeVeniceCluster(); + String storeName = "test-store"; Consumer paramsConsumer = params -> params.setBlobTransferEnabled(true); setUpBatchStore(cluster, storeName, paramsConsumer, properties -> {}, true); - VeniceServerWrapper server1 = cluster.getVeniceServers().get(0); - VeniceServerWrapper server2 = cluster.getVeniceServers().get(1); + VeniceServerWrapper server1 = cluster.getVeniceServerByPort(server1Port); + VeniceServerWrapper server2 = cluster.getVeniceServerByPort(server2Port); // verify the snapshot is generated for both servers after the job is done for (int partitionId = 0; partitionId < PARTITION_COUNT; partitionId++) { @@ -97,7 +95,7 @@ public void testBlobP2PTransferAmongServersForBatchStore() throws Exception { Files.exists(Paths.get(RocksDBUtils.composeSnapshotDir(path1 + "/rocksdb", storeName + "_v1", partitionId)))); } - cluster.stopAndRestartVeniceServer(server1.getPort()); + cluster.stopAndRestartVeniceServer(server1Port); TestUtils.waitForNonDeterministicAssertion(2, TimeUnit.MINUTES, () -> { Assert.assertTrue(server1.isRunning()); }); @@ -121,6 +119,11 @@ public void testBlobP2PTransferAmongServersForBatchStore() throws Exception { File file = new File(RocksDBUtils.composePartitionDbDir(path1 + "/rocksdb", storeName + "_v1", partitionId)); Boolean fileExisted = Files.exists(file.toPath()); Assert.assertTrue(fileExisted); + // ensure the snapshot file is not generated + File snapshotFile = + new File(RocksDBUtils.composeSnapshotDir(path1 + "/rocksdb", storeName + "_v1", partitionId)); + Boolean snapshotFileExisted = Files.exists(snapshotFile.toPath()); + Assert.assertFalse(snapshotFileExisted); } }); @@ -141,12 +144,12 @@ public void testBlobP2PTransferAmongServersForBatchStore() throws Exception { */ @Test(singleThreaded = true, timeOut = 180000) public void testBlobTransferThrowExceptionIfSnapshotNotExisted() throws Exception { + cluster = initializeVeniceCluster(); String storeName = "test-store-snapshot-not-existed"; Consumer paramsConsumer = params -> params.setBlobTransferEnabled(true); setUpBatchStore(cluster, storeName, paramsConsumer, properties -> {}, true); - VeniceServerWrapper server1 = cluster.getVeniceServers().get(0); - + VeniceServerWrapper server1 = cluster.getVeniceServerByPort(server1Port); // verify the snapshot is generated for both servers after the job is done for (int i = 0; i < PARTITION_COUNT; i++) { String snapshotPath1 = RocksDBUtils.composeSnapshotDir(path1 + "/rocksdb", storeName + "_v1", i); @@ -176,7 +179,7 @@ public void testBlobTransferThrowExceptionIfSnapshotNotExisted() throws Exceptio } } - cluster.stopAndRestartVeniceServer(server1.getPort()); + cluster.stopAndRestartVeniceServer(server1Port); TestUtils.waitForNonDeterministicAssertion(2, TimeUnit.MINUTES, () -> { Assert.assertTrue(server1.isRunning()); }); @@ -207,7 +210,81 @@ public void testBlobTransferThrowExceptionIfSnapshotNotExisted() throws Exceptio }); } + /** + * Test when the format of the rocksdb is different between two servers, + * the blob transfer should throw an exception and return a 404 error. + */ + @Test(singleThreaded = true, timeOut = 180000) + public void testBlobTransferThrowExceptionIfTableFormatNotMatch() throws Exception { + cluster = initializeVeniceCluster(false); + + String storeName = "test-store-format-not-match"; + Consumer paramsConsumer = params -> params.setBlobTransferEnabled(true); + setUpBatchStore(cluster, storeName, paramsConsumer, properties -> {}, true); + + VeniceServerWrapper server1 = cluster.getVeniceServerByPort(server1Port); + + // verify the snapshot is generated for both servers after the job is done + for (int partitionId = 0; partitionId < PARTITION_COUNT; partitionId++) { + String snapshotPath1 = RocksDBUtils.composeSnapshotDir(path1 + "/rocksdb", storeName + "_v1", partitionId); + Assert.assertTrue(Files.exists(Paths.get(snapshotPath1))); + String snapshotPath2 = RocksDBUtils.composeSnapshotDir(path2 + "/rocksdb", storeName + "_v1", partitionId); + Assert.assertTrue(Files.exists(Paths.get(snapshotPath2))); + } + + // cleanup and restart server 1 + FileUtils.deleteDirectory( + new File(RocksDBUtils.composePartitionDbDir(path1 + "/rocksdb", storeName + "_v1", METADATA_PARTITION_ID))); + for (int partitionId = 0; partitionId < PARTITION_COUNT; partitionId++) { + FileUtils.deleteDirectory( + new File(RocksDBUtils.composePartitionDbDir(path1 + "/rocksdb", storeName + "_v1", partitionId))); + // both partition db and snapshot should be deleted + Assert.assertFalse( + Files.exists( + Paths.get(RocksDBUtils.composePartitionDbDir(path1 + "/rocksdb", storeName + "_v1", partitionId)))); + Assert.assertFalse( + Files.exists(Paths.get(RocksDBUtils.composeSnapshotDir(path1 + "/rocksdb", storeName + "_v1", partitionId)))); + } + + cluster.stopAndRestartVeniceServer(server1Port); + TestUtils.waitForNonDeterministicAssertion(2, TimeUnit.MINUTES, () -> { + Assert.assertTrue(server1.isRunning()); + }); + + // wait for server 1 + cluster.getVeniceControllers().forEach(controller -> { + TestUtils.waitForNonDeterministicAssertion(2, TimeUnit.MINUTES, () -> { + Assert.assertEquals( + controller.getController() + .getVeniceControllerService() + .getVeniceHelixAdmin() + .getAllStoreStatuses(cluster.getClusterName()) + .get(storeName), + FULLLY_REPLICATED.toString()); + }); + }); + + TestUtils.waitForNonDeterministicAssertion(2, TimeUnit.MINUTES, () -> { + for (int partitionId = 0; partitionId < PARTITION_COUNT; partitionId++) { + File file = new File(RocksDBUtils.composePartitionDbDir(path1 + "/rocksdb", storeName + "_v1", partitionId)); + Boolean fileExisted = Files.exists(file.toPath()); + Assert.assertTrue(fileExisted); + // ensure that the snapshot file is generated as it is re-ingested, not from server2 file transfer + File snapshotFile = + new File(RocksDBUtils.composeSnapshotDir(path1 + "/rocksdb", storeName + "_v1", partitionId)); + Boolean snapshotFileExisted = Files.exists(snapshotFile.toPath()); + Assert.assertTrue(snapshotFileExisted); + } + }); + } + public VeniceClusterWrapper initializeVeniceCluster() { + return initializeVeniceCluster(true); + } + + public VeniceClusterWrapper initializeVeniceCluster(boolean sameRocksDBFormat) { + server1Port = -1; + server2Port = -1; path1 = Utils.getTempDataDirectory().getAbsolutePath(); path2 = Utils.getTempDataDirectory().getAbsolutePath(); @@ -231,6 +308,8 @@ public VeniceClusterWrapper initializeVeniceCluster() { serverProperties.setProperty(ConfigKeys.DAVINCI_P2P_BLOB_TRANSFER_CLIENT_PORT, String.valueOf(port2)); serverProperties.setProperty(ConfigKeys.BLOB_TRANSFER_MANAGER_ENABLED, "true"); veniceClusterWrapper.addVeniceServer(new Properties(), serverProperties); + // get the first port id for finding first server. + server1Port = veniceClusterWrapper.getVeniceServers().get(0).getPort(); // add second server serverProperties.setProperty(ConfigKeys.DATA_BASE_PATH, path2); @@ -238,7 +317,20 @@ public VeniceClusterWrapper initializeVeniceCluster() { serverProperties.setProperty(ConfigKeys.BLOB_TRANSFER_MANAGER_ENABLED, "true"); serverProperties.setProperty(ConfigKeys.DAVINCI_P2P_BLOB_TRANSFER_SERVER_PORT, String.valueOf(port2)); serverProperties.setProperty(ConfigKeys.DAVINCI_P2P_BLOB_TRANSFER_CLIENT_PORT, String.valueOf(port1)); + + if (!sameRocksDBFormat) { + // the second server use PLAIN_TABLE_FORMAT + serverProperties.setProperty(ROCKSDB_PLAIN_TABLE_FORMAT_ENABLED, "true"); + } + veniceClusterWrapper.addVeniceServer(new Properties(), serverProperties); + // get the second port num for finding second server, + // because the order of servers is not guaranteed, need to exclude the first server. + for (VeniceServerWrapper server: veniceClusterWrapper.getVeniceServers()) { + if (server.getPort() != server1Port) { + server2Port = server.getPort(); + } + } Properties routerProperties = new Properties(); routerProperties.put(ConfigKeys.ROUTER_CLIENT_DECOMPRESSION_ENABLED, "true"); @@ -288,6 +380,8 @@ private static void runVPJ(Properties vpjProperties, int expectedVersionNumber, @Test(singleThreaded = true, timeOut = 180000) public void testBlobP2PTransferAmongServersForHybridStore() throws Exception { + cluster = initializeVeniceCluster(); + ControllerClient controllerClient = new ControllerClient(cluster.getClusterName(), cluster.getAllControllersURLs()); // prepare hybrid store. String storeName = "test-store-hybrid"; @@ -303,8 +397,8 @@ public void testBlobP2PTransferAmongServersForHybridStore() throws Exception { TestUtils.assertCommand( controllerClient.sendEmptyPushAndWait(storeName, Utils.getUniqueString("empty-hybrid-push"), 1L, 120000)); - VeniceServerWrapper server1 = cluster.getVeniceServers().get(0); - VeniceServerWrapper server2 = cluster.getVeniceServers().get(1); + VeniceServerWrapper server1 = cluster.getVeniceServerByPort(server1Port); + VeniceServerWrapper server2 = cluster.getVeniceServerByPort(server2Port); // offset record should be same after the empty push TestUtils.waitForNonDeterministicAssertion(2, TimeUnit.MINUTES, () -> { @@ -318,7 +412,7 @@ public void testBlobP2PTransferAmongServersForHybridStore() throws Exception { }); // cleanup and stop server 1 - cluster.stopVeniceServer(server1.getPort()); + cluster.stopVeniceServer(server1Port); FileUtils.deleteDirectory( new File(RocksDBUtils.composePartitionDbDir(path1 + "/rocksdb", storeName + "_v1", METADATA_PARTITION_ID))); for (int partitionId = 0; partitionId < PARTITION_COUNT; partitionId++) { @@ -342,9 +436,9 @@ public void testBlobP2PTransferAmongServersForHybridStore() throws Exception { veniceProducer.stop(); } + cluster.restartVeniceServer(server1.getPort()); // restart server 1 TestUtils.waitForNonDeterministicAssertion(2, TimeUnit.MINUTES, () -> { - cluster.restartVeniceServer(server1.getPort()); Assert.assertTrue(server1.isRunning()); }); @@ -366,6 +460,11 @@ public void testBlobP2PTransferAmongServersForHybridStore() throws Exception { File file = new File(RocksDBUtils.composePartitionDbDir(path1 + "/rocksdb", storeName + "_v1", partitionId)); Boolean fileExisted = Files.exists(file.toPath()); Assert.assertTrue(fileExisted); + // ensure that the snapshot is not generated. + File snapshotFile = + new File(RocksDBUtils.composeSnapshotDir(path1 + "/rocksdb", storeName + "_v1", partitionId)); + Boolean snapshotFileExisted = Files.exists(snapshotFile.toPath()); + Assert.assertFalse(snapshotFileExisted); } }); diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientRecordTransformerTest.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientRecordTransformerTest.java index fcadab46499..52fc30e12c9 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientRecordTransformerTest.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DaVinciClientRecordTransformerTest.java @@ -114,8 +114,14 @@ public void testRecordTransformer(DaVinciConfig clientConfig) throws Exception { VeniceRouterWrapper.CLUSTER_DISCOVERY_D2_SERVICE_NAME, metricsRepository, backendConfig)) { + DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( - (storeVersion) -> new TestStringRecordTransformer(storeVersion, true), + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> new TestStringRecordTransformer( + storeVersion, + keySchema, + inputValueSchema, + outputValueSchema, + true), String.class, Schema.create(Schema.Type.STRING)); clientConfig.setRecordTransformerConfig(recordTransformerConfig); @@ -156,7 +162,12 @@ public void testTypeChangeRecordTransformer(DaVinciConfig clientConfig) throws E metricsRepository, backendConfig)) { DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( - (storeVersion) -> new TestIntToStringRecordTransformer(storeVersion, true), + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> new TestIntToStringRecordTransformer( + storeVersion, + keySchema, + inputValueSchema, + outputValueSchema, + true), String.class, Schema.create(Schema.Type.STRING)); clientConfig.setRecordTransformerConfig(recordTransformerConfig); @@ -199,10 +210,14 @@ public void testRecordTransformerOnRecovery(DaVinciConfig clientConfig) throws E metricsRepository, backendConfig)) { - TestStringRecordTransformer recordTransformer = new TestStringRecordTransformer(1, true); + Schema myKeySchema = Schema.create(Schema.Type.INT); + Schema myValueSchema = Schema.create(Schema.Type.STRING); + + TestStringRecordTransformer recordTransformer = + new TestStringRecordTransformer(1, myKeySchema, myValueSchema, myValueSchema, true); DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( - (storeVersion) -> recordTransformer, + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> recordTransformer, String.class, Schema.create(Schema.Type.STRING)); clientConfig.setRecordTransformerConfig(recordTransformerConfig); @@ -274,10 +289,14 @@ public void testRecordTransformerChunking(DaVinciConfig clientConfig) throws Exc metricsRepository, backendConfig)) { - TestStringRecordTransformer recordTransformer = new TestStringRecordTransformer(1, true); + Schema myKeySchema = Schema.create(Schema.Type.INT); + Schema myValueSchema = Schema.create(Schema.Type.STRING); + + TestStringRecordTransformer recordTransformer = + new TestStringRecordTransformer(1, myKeySchema, myValueSchema, myValueSchema, true); DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( - (storeVersion) -> recordTransformer, + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> recordTransformer, String.class, Schema.create(Schema.Type.STRING)); clientConfig.setRecordTransformerConfig(recordTransformerConfig); @@ -334,17 +353,23 @@ public void testRecordTransformerWithEmptyDaVinci(DaVinciConfig clientConfig) th VeniceProperties backendConfig = buildRecordTransformerBackendConfig(pushStatusStoreEnabled); MetricsRepository metricsRepository = new MetricsRepository(); - TestStringRecordTransformer recordTransformer = new TestStringRecordTransformer(1, false); + Schema myKeySchema = Schema.create(Schema.Type.INT); + Schema myValueSchema = Schema.create(Schema.Type.STRING); + + TestStringRecordTransformer recordTransformer = + new TestStringRecordTransformer(1, myKeySchema, myValueSchema, myValueSchema, false); + + DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> recordTransformer, + String.class, + Schema.create(Schema.Type.STRING)); + clientConfig.setRecordTransformerConfig(recordTransformerConfig); try (CachingDaVinciClientFactory factory = new CachingDaVinciClientFactory( d2Client, VeniceRouterWrapper.CLUSTER_DISCOVERY_D2_SERVICE_NAME, metricsRepository, backendConfig)) { - DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( - (storeVersion) -> recordTransformer, - String.class, - Schema.create(Schema.Type.STRING)); clientConfig.setRecordTransformerConfig(recordTransformerConfig); DaVinciClient clientWithRecordTransformer = @@ -377,17 +402,23 @@ public void testSkipResultRecordTransformer(DaVinciConfig clientConfig) throws E VeniceProperties backendConfig = buildRecordTransformerBackendConfig(pushStatusStoreEnabled); MetricsRepository metricsRepository = new MetricsRepository(); - TestSkipResultRecordTransformer recordTransformer = new TestSkipResultRecordTransformer(1, true); + Schema myKeySchema = Schema.create(Schema.Type.INT); + Schema myValueSchema = Schema.create(Schema.Type.STRING); + + TestSkipResultRecordTransformer recordTransformer = + new TestSkipResultRecordTransformer(1, myKeySchema, myValueSchema, myValueSchema, true); + + DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> recordTransformer, + String.class, + Schema.create(Schema.Type.STRING)); + clientConfig.setRecordTransformerConfig(recordTransformerConfig); try (CachingDaVinciClientFactory factory = new CachingDaVinciClientFactory( d2Client, VeniceRouterWrapper.CLUSTER_DISCOVERY_D2_SERVICE_NAME, metricsRepository, backendConfig)) { - DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( - (storeVersion) -> recordTransformer, - String.class, - Schema.create(Schema.Type.STRING)); clientConfig.setRecordTransformerConfig(recordTransformerConfig); DaVinciClient clientWithRecordTransformer = @@ -420,15 +451,22 @@ public void testUnchangedResultRecordTransformer(DaVinciConfig clientConfig) thr VeniceProperties backendConfig = buildRecordTransformerBackendConfig(pushStatusStoreEnabled); MetricsRepository metricsRepository = new MetricsRepository(); + DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> new TestUnchangedResultRecordTransformer( + storeVersion, + keySchema, + inputValueSchema, + outputValueSchema, + true), + String.class, + Schema.create(Schema.Type.STRING)); + clientConfig.setRecordTransformerConfig(recordTransformerConfig); + try (CachingDaVinciClientFactory factory = new CachingDaVinciClientFactory( d2Client, VeniceRouterWrapper.CLUSTER_DISCOVERY_D2_SERVICE_NAME, metricsRepository, backendConfig)) { - DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( - (storeVersion) -> new TestUnchangedResultRecordTransformer(storeVersion, true), - String.class, - Schema.create(Schema.Type.STRING)); clientConfig.setRecordTransformerConfig(recordTransformerConfig); DaVinciClient clientWithRecordTransformer = diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DuckDBDaVinciRecordTransformerIntegrationTest.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DuckDBDaVinciRecordTransformerIntegrationTest.java new file mode 100644 index 00000000000..71b1acb0f24 --- /dev/null +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/DuckDBDaVinciRecordTransformerIntegrationTest.java @@ -0,0 +1,258 @@ +package com.linkedin.venice.endToEnd; + +import static com.linkedin.venice.ConfigKeys.CLIENT_SYSTEM_STORE_REPOSITORY_REFRESH_INTERVAL_SECONDS; +import static com.linkedin.venice.ConfigKeys.CLIENT_USE_SYSTEM_STORE_REPOSITORY; +import static com.linkedin.venice.ConfigKeys.DATA_BASE_PATH; +import static com.linkedin.venice.ConfigKeys.DAVINCI_PUSH_STATUS_CHECK_INTERVAL_IN_MS; +import static com.linkedin.venice.ConfigKeys.DAVINCI_PUSH_STATUS_SCAN_INTERVAL_IN_SECONDS; +import static com.linkedin.venice.ConfigKeys.DA_VINCI_CURRENT_VERSION_BOOTSTRAPPING_SPEEDUP_ENABLED; +import static com.linkedin.venice.ConfigKeys.PERSISTENCE_TYPE; +import static com.linkedin.venice.ConfigKeys.PUSH_STATUS_STORE_ENABLED; +import static com.linkedin.venice.ConfigKeys.SERVER_PROMOTION_TO_LEADER_REPLICA_DELAY_SECONDS; +import static com.linkedin.venice.meta.PersistenceType.ROCKS_DB; +import static com.linkedin.venice.utils.IntegrationTestPushUtils.createStoreForJob; +import static com.linkedin.venice.utils.IntegrationTestPushUtils.defaultVPJProps; +import static com.linkedin.venice.utils.TestWriteUtils.DEFAULT_USER_DATA_RECORD_COUNT; +import static com.linkedin.venice.utils.TestWriteUtils.NAME_RECORD_V1_SCHEMA; +import static com.linkedin.venice.utils.TestWriteUtils.SINGLE_FIELD_RECORD_SCHEMA; +import static com.linkedin.venice.utils.TestWriteUtils.getTempDataDirectory; +import static com.linkedin.venice.utils.TestWriteUtils.writeSimpleAvroFile; +import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_KEY_FIELD_PROP; +import static com.linkedin.venice.vpj.VenicePushJobConstants.DEFAULT_VALUE_FIELD_PROP; +import static com.linkedin.venice.vpj.VenicePushJobConstants.VENICE_STORE_NAME_PROP; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; + +import com.linkedin.d2.balancer.D2Client; +import com.linkedin.d2.balancer.D2ClientBuilder; +import com.linkedin.davinci.client.DaVinciClient; +import com.linkedin.davinci.client.DaVinciConfig; +import com.linkedin.davinci.client.DaVinciRecordTransformerConfig; +import com.linkedin.davinci.client.factory.CachingDaVinciClientFactory; +import com.linkedin.venice.D2.D2ClientUtils; +import com.linkedin.venice.compression.CompressionStrategy; +import com.linkedin.venice.controllerapi.ControllerClient; +import com.linkedin.venice.controllerapi.SchemaResponse; +import com.linkedin.venice.controllerapi.UpdateStoreQueryParams; +import com.linkedin.venice.duckdb.DuckDBDaVinciRecordTransformer; +import com.linkedin.venice.integration.utils.ServiceFactory; +import com.linkedin.venice.integration.utils.VeniceClusterWrapper; +import com.linkedin.venice.integration.utils.VeniceRouterWrapper; +import com.linkedin.venice.utils.PropertyBuilder; +import com.linkedin.venice.utils.PushInputSchemaBuilder; +import com.linkedin.venice.utils.TestUtils; +import com.linkedin.venice.utils.TestWriteUtils; +import com.linkedin.venice.utils.Time; +import com.linkedin.venice.utils.Utils; +import com.linkedin.venice.utils.VeniceProperties; +import io.tehuti.metrics.MetricsRepository; +import java.io.File; +import java.io.IOException; +import java.sql.Connection; +import java.sql.DriverManager; +import java.sql.ResultSet; +import java.sql.SQLException; +import java.sql.Statement; +import java.util.Collections; +import java.util.Properties; +import java.util.Set; +import java.util.concurrent.TimeUnit; +import java.util.function.Consumer; +import org.apache.avro.Schema; +import org.apache.avro.generic.GenericData; +import org.apache.avro.generic.GenericRecord; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class DuckDBDaVinciRecordTransformerIntegrationTest { + private static final Logger LOGGER = LogManager.getLogger(DaVinciClientRecordTransformerTest.class); + private static final int TEST_TIMEOUT = 3 * Time.MS_PER_MINUTE; + private VeniceClusterWrapper cluster; + private D2Client d2Client; + + @BeforeClass + public void setUp() { + Utils.thisIsLocalhost(); + Properties clusterConfig = new Properties(); + clusterConfig.put(SERVER_PROMOTION_TO_LEADER_REPLICA_DELAY_SECONDS, 1L); + clusterConfig.put(PUSH_STATUS_STORE_ENABLED, true); + clusterConfig.put(DAVINCI_PUSH_STATUS_SCAN_INTERVAL_IN_SECONDS, 3); + cluster = ServiceFactory.getVeniceCluster(1, 2, 1, 2, 100, false, false, clusterConfig); + d2Client = new D2ClientBuilder().setZkHosts(cluster.getZk().getAddress()) + .setZkSessionTimeout(3, TimeUnit.SECONDS) + .setZkStartupTimeout(3, TimeUnit.SECONDS) + .build(); + D2ClientUtils.startClient(d2Client); + } + + @AfterClass + public void cleanUp() { + if (d2Client != null) { + D2ClientUtils.shutdownClient(d2Client); + } + Utils.closeQuietlyWithErrorLogged(cluster); + } + + @BeforeMethod + @AfterClass + public void deleteClassHash() { + int storeVersion = 1; + File file = new File(String.format("./classHash-%d.txt", storeVersion)); + if (file.exists()) { + assertTrue(file.delete()); + } + } + + /** + * This test works on mac but fails on the CI, likely due to: https://github.com/duckdb/duckdb-java/issues/14 + * + * There is a fix merged for this, but it has not been released yet. + * + * TODO: Re-enable once we can depend on a clean release. + */ + @Test(timeOut = TEST_TIMEOUT) + public void testRecordTransformer() throws Exception { + DaVinciConfig clientConfig = new DaVinciConfig(); + + File tmpDir = Utils.getTempDataDirectory(); + String storeName = Utils.getUniqueString("test_store"); + boolean pushStatusStoreEnabled = false; + boolean chunkingEnabled = false; + CompressionStrategy compressionStrategy = CompressionStrategy.NO_OP; + + setUpStore(storeName, pushStatusStoreEnabled, chunkingEnabled, compressionStrategy); + + VeniceProperties backendConfig = buildRecordTransformerBackendConfig(pushStatusStoreEnabled); + MetricsRepository metricsRepository = new MetricsRepository(); + String duckDBUrl = "jdbc:duckdb:" + tmpDir.getAbsolutePath() + "/my_database.duckdb"; + + try (CachingDaVinciClientFactory factory = new CachingDaVinciClientFactory( + d2Client, + VeniceRouterWrapper.CLUSTER_DISCOVERY_D2_SERVICE_NAME, + metricsRepository, + backendConfig)) { + Set columnsToProject = Collections.emptySet(); + DaVinciRecordTransformerConfig recordTransformerConfig = new DaVinciRecordTransformerConfig( + (storeVersion, keySchema, inputValueSchema, outputValueSchema) -> new DuckDBDaVinciRecordTransformer( + storeVersion, + keySchema, + inputValueSchema, + outputValueSchema, + false, + tmpDir.getAbsolutePath(), + storeName, + columnsToProject), + GenericRecord.class, + NAME_RECORD_V1_SCHEMA); + clientConfig.setRecordTransformerConfig(recordTransformerConfig); + + DaVinciClient clientWithRecordTransformer = + factory.getAndStartGenericAvroClient(storeName, clientConfig); + + clientWithRecordTransformer.subscribeAll().get(); + + assertRowCount(duckDBUrl, "subscribeAll() finishes!"); + + clientWithRecordTransformer.unsubscribeAll(); + } + + assertRowCount(duckDBUrl, "DVC gets closed!"); + } + + private void assertRowCount(String duckDBUrl, String assertionErrorMsg) throws SQLException { + try (Connection connection = DriverManager.getConnection(duckDBUrl); + Statement statement = connection.createStatement(); + ResultSet rs = statement.executeQuery("SELECT count(*) FROM current_version")) { + assertTrue(rs.next()); + int rowCount = rs.getInt(1); + assertEquals( + rowCount, + DEFAULT_USER_DATA_RECORD_COUNT, + "The DB should contain " + DEFAULT_USER_DATA_RECORD_COUNT + " right after " + assertionErrorMsg); + } + } + + protected void setUpStore( + String storeName, + boolean useDVCPushStatusStore, + boolean chunkingEnabled, + CompressionStrategy compressionStrategy) throws IOException { + + File inputDir = getTempDataDirectory(); + Consumer paramsConsumer = params -> {}; + Consumer propertiesConsumer = properties -> {}; + Schema pushRecordSchema = new PushInputSchemaBuilder().setKeySchema(SINGLE_FIELD_RECORD_SCHEMA) + .setValueSchema(NAME_RECORD_V1_SCHEMA) + .build(); + String firstName = "first_name_"; + String lastName = "last_name_"; + Schema valueSchema = writeSimpleAvroFile(inputDir, pushRecordSchema, i -> { + GenericRecord keyValueRecord = new GenericData.Record(pushRecordSchema); + GenericRecord key = new GenericData.Record(SINGLE_FIELD_RECORD_SCHEMA); + key.put("key", i.toString()); + keyValueRecord.put(DEFAULT_KEY_FIELD_PROP, key); + GenericRecord valueRecord = new GenericData.Record(NAME_RECORD_V1_SCHEMA); + valueRecord.put("firstName", firstName + i); + valueRecord.put("lastName", lastName + i); + keyValueRecord.put(DEFAULT_VALUE_FIELD_PROP, valueRecord); // Value + return keyValueRecord; + }); + String keySchemaStr = valueSchema.getField(DEFAULT_KEY_FIELD_PROP).schema().toString(); + + // Setup VPJ job properties. + String inputDirPath = "file://" + inputDir.getAbsolutePath(); + Properties vpjProperties = defaultVPJProps(cluster, inputDirPath, storeName); + propertiesConsumer.accept(vpjProperties); + // Create & update store for test. + final int numPartitions = 3; + UpdateStoreQueryParams params = new UpdateStoreQueryParams().setPartitionCount(numPartitions) + .setChunkingEnabled(chunkingEnabled) + .setCompressionStrategy(compressionStrategy); + + paramsConsumer.accept(params); + + try (ControllerClient controllerClient = + createStoreForJob(cluster, keySchemaStr, NAME_RECORD_V1_SCHEMA.toString(), vpjProperties)) { + cluster.createMetaSystemStore(storeName); + if (useDVCPushStatusStore) { + cluster.createPushStatusSystemStore(storeName); + } + TestUtils.assertCommand(controllerClient.updateStore(storeName, params)); + SchemaResponse schemaResponse = controllerClient.addValueSchema(storeName, NAME_RECORD_V1_SCHEMA.toString()); + assertFalse(schemaResponse.isError()); + runVPJ(vpjProperties, 1, cluster); + } + } + + private static void runVPJ(Properties vpjProperties, int expectedVersionNumber, VeniceClusterWrapper cluster) { + long vpjStart = System.currentTimeMillis(); + String jobName = Utils.getUniqueString("batch-job-" + expectedVersionNumber); + TestWriteUtils.runPushJob(jobName, vpjProperties); + String storeName = (String) vpjProperties.get(VENICE_STORE_NAME_PROP); + cluster.waitVersion(storeName, expectedVersionNumber); + LOGGER.info("**TIME** VPJ" + expectedVersionNumber + " takes " + (System.currentTimeMillis() - vpjStart)); + } + + public VeniceProperties buildRecordTransformerBackendConfig(boolean pushStatusStoreEnabled) { + String baseDataPath = Utils.getTempDataDirectory().getAbsolutePath(); + PropertyBuilder backendPropertyBuilder = new PropertyBuilder().put(CLIENT_USE_SYSTEM_STORE_REPOSITORY, true) + .put(CLIENT_SYSTEM_STORE_REPOSITORY_REFRESH_INTERVAL_SECONDS, 1) + .put(DATA_BASE_PATH, baseDataPath) + .put(PERSISTENCE_TYPE, ROCKS_DB) + .put(DA_VINCI_CURRENT_VERSION_BOOTSTRAPPING_SPEEDUP_ENABLED, true) + .put(PUSH_STATUS_STORE_ENABLED, pushStatusStoreEnabled) + .put(DAVINCI_PUSH_STATUS_CHECK_INTERVAL_IN_MS, 1000); + + if (pushStatusStoreEnabled) { + backendPropertyBuilder.put(PUSH_STATUS_STORE_ENABLED, true).put(DAVINCI_PUSH_STATUS_CHECK_INTERVAL_IN_MS, 1000); + } + + return backendPropertyBuilder.build(); + } +} diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/MetaSystemStoreTest.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/MetaSystemStoreTest.java index cbb4bcd1d42..7dc55dc7b34 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/MetaSystemStoreTest.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/MetaSystemStoreTest.java @@ -13,7 +13,6 @@ import static org.testng.Assert.expectThrows; import static org.testng.Assert.fail; -import com.google.gson.Gson; import com.linkedin.d2.balancer.D2Client; import com.linkedin.davinci.repository.NativeMetadataRepository; import com.linkedin.davinci.repository.RequestBasedMetaRepository; @@ -40,6 +39,7 @@ import com.linkedin.venice.meta.Store; import com.linkedin.venice.meta.Version; import com.linkedin.venice.meta.VersionStatus; +import com.linkedin.venice.meta.ZKStore; import com.linkedin.venice.schema.SchemaEntry; import com.linkedin.venice.serialization.avro.AvroProtocolDefinition; import com.linkedin.venice.system.store.MetaStoreDataType; @@ -289,7 +289,6 @@ public void testThinClientMetaStoreBasedRepository() throws InterruptedException } } - // TODO PRANAV TEST @Test(timeOut = 120 * Time.MS_PER_SECOND) public void testRequestBasedMetaStoreBasedRepository() throws InterruptedException { String regularVeniceStoreName = Utils.getUniqueString("venice_store"); @@ -433,15 +432,13 @@ private void verifyRepository(NativeMetadataRepository nativeMetadataRepository, expectThrows(VeniceNoStoreException.class, () -> nativeMetadataRepository.getStoreOrThrow("Non-existing-store")); expectThrows(VeniceNoStoreException.class, () -> nativeMetadataRepository.subscribe("Non-existing-store")); nativeMetadataRepository.subscribe(regularVeniceStoreName); - Store store = new ReadOnlyStore(nativeMetadataRepository.getStore(regularVeniceStoreName)); - Store controllerStore = new ReadOnlyStore( - veniceLocalCluster.getLeaderVeniceController().getVeniceAdmin().getStore(clusterName, regularVeniceStoreName)); - // TODO PRANAV this is failing, stores are not exaclty the same - // Strings are CharSeqs in actual store - System.out.println(new Gson().toJson(store)); - System.out.println("======================="); - System.out.println(new Gson().toJson(controllerStore)); - // assertEquals(store, controllerStore); + Store store = normalizeStore(new ReadOnlyStore(nativeMetadataRepository.getStore(regularVeniceStoreName))); + Store controllerStore = normalizeStore( + new ReadOnlyStore( + veniceLocalCluster.getLeaderVeniceController() + .getVeniceAdmin() + .getStore(clusterName, regularVeniceStoreName))); + assertEquals(store.toString(), controllerStore.toString()); SchemaEntry keySchema = nativeMetadataRepository.getKeySchema(regularVeniceStoreName); SchemaEntry controllerKeySchema = veniceLocalCluster.getLeaderVeniceController() .getVeniceAdmin() @@ -491,6 +488,10 @@ private void verifyRepository(NativeMetadataRepository nativeMetadataRepository, }); } + private Store normalizeStore(ReadOnlyStore store) { + return new ReadOnlyStore(new ZKStore(store.cloneStoreProperties())); + } + private void createStoreAndMaterializeMetaSystemStore(String storeName) { createStoreAndMaterializeMetaSystemStore(storeName, VALUE_SCHEMA_1); } @@ -498,7 +499,6 @@ private void createStoreAndMaterializeMetaSystemStore(String storeName) { private void createStoreAndMaterializeMetaSystemStore(String storeName, String valueSchema) { // Verify and create Venice regular store if it doesn't exist. if (parentControllerClient.getStore(storeName).getStore() == null) { - // Flaky? NewStoreResponse resp = parentControllerClient.createNewStore(storeName, "test_owner", INT_KEY_SCHEMA, valueSchema); if (resp.isError()) { diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/PartialUpdateTest.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/PartialUpdateTest.java index 901245c4a73..ab31cccda4f 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/PartialUpdateTest.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/PartialUpdateTest.java @@ -174,6 +174,7 @@ public void setUp() { Boolean.toString(isAAWCParallelProcessingEnabled())); Properties controllerProps = new Properties(); controllerProps.put(ConfigKeys.CONTROLLER_AUTO_MATERIALIZE_META_SYSTEM_STORE, false); + controllerProps.put(ConfigKeys.PARTICIPANT_MESSAGE_STORE_ENABLED, false); this.multiRegionMultiClusterWrapper = ServiceFactory.getVeniceTwoLayerMultiRegionMultiClusterWrapper( NUMBER_OF_CHILD_DATACENTERS, NUMBER_OF_CLUSTERS, @@ -356,7 +357,7 @@ public void testIncrementalPushPartialUpdateClassicFormat() throws IOException { TestUtils.waitForNonDeterministicPushCompletion( Version.composeKafkaTopic(storeName, 1), parentControllerClient, - 30, + 60, TimeUnit.SECONDS); // VPJ push diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestActiveActiveReplicationForIncPush.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestActiveActiveReplicationForIncPush.java index b50bc298565..63fb50cbaf5 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestActiveActiveReplicationForIncPush.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestActiveActiveReplicationForIncPush.java @@ -21,6 +21,7 @@ import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertNotNull; import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; import com.linkedin.venice.controllerapi.ControllerClient; import com.linkedin.venice.controllerapi.StoreResponse; @@ -37,8 +38,11 @@ import com.linkedin.venice.meta.Version; import com.linkedin.venice.pubsub.PubSubTopicPartitionImpl; import com.linkedin.venice.pubsub.PubSubTopicRepository; +import com.linkedin.venice.pubsub.api.PubSubTopic; import com.linkedin.venice.pubsub.api.PubSubTopicPartition; import com.linkedin.venice.pubsub.manager.TopicManager; +import com.linkedin.venice.samza.VeniceSystemFactory; +import com.linkedin.venice.samza.VeniceSystemProducer; import com.linkedin.venice.utils.IntegrationTestPushUtils; import com.linkedin.venice.utils.TestUtils; import com.linkedin.venice.utils.TestWriteUtils; @@ -47,6 +51,7 @@ import java.io.File; import java.util.ArrayList; import java.util.List; +import java.util.Map; import java.util.Optional; import java.util.Properties; import java.util.concurrent.CompletableFuture; @@ -55,6 +60,7 @@ import org.apache.avro.Schema; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import org.apache.samza.config.MapConfig; import org.testng.Assert; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; @@ -134,6 +140,135 @@ public void testAAReplicationForIncPush() throws Exception { testAAReplicationForIncPush(false); } + /** + * This test reproduces an issue where the real-time topic partition count did not match the hybrid version + * partition count under the following scenario: + * + * 1. Create a store with 1 partition. + * 2. Perform a batch push, resulting in a batch version with 1 partition. + * 3. Update the store to have 3 partitions and convert it into a hybrid store. + * 4. Start real-time writes using push type {@link com.linkedin.venice.meta.Version.PushType#STREAM}. + * 5. Perform a full push, which creates a hybrid version with 3 partitions. This push results in an error + * because, after the topic switch to real-time consumers, partitions 1 and 2 of the real-time topic cannot + * be found, as it has only 1 partition (partition: 0). + * + * The root cause of the issue lies in step 4, where the real-time topic was created if it did not already exist. + * The partition count for the real-time topic was derived from the largest existing version, which in this case + * was the batch version with 1 partition. This caused the real-time topic to have incorrect partition count (1 + * instead of 3). + * + * To resolve this issue: + * - STREAM push type is no longer allowed if there is no online hybrid version. + * - If there is an online hybrid version, it is safe to assume that the real-time topic partition count matches + * the hybrid version partition count. + * - The real-time topic is no longer created if it does not exist as part of the `requestTopicForPushing` method. + */ + @Test(timeOut = TEST_TIMEOUT) + public void testRealTimeTopicPartitionCountMatchesHybridVersion() throws Exception { + File inputDirBatch = getTempDataDirectory(); + String parentControllerUrls = multiRegionMultiClusterWrapper.getControllerConnectString(); + String inputDirPathBatch = "file:" + inputDirBatch.getAbsolutePath(); + try (ControllerClient parentControllerClient = new ControllerClient(clusterName, parentControllerUrls)) { + String storeName = Utils.getUniqueString("store"); + Properties propsBatch = + IntegrationTestPushUtils.defaultVPJProps(multiRegionMultiClusterWrapper, inputDirPathBatch, storeName); + propsBatch.put(SEND_CONTROL_MESSAGES_DIRECTLY, true); + Schema recordSchema = TestWriteUtils.writeSimpleAvroFileWithStringToStringSchema(inputDirBatch); + String keySchemaStr = recordSchema.getField(DEFAULT_KEY_FIELD_PROP).schema().toString(); + String valueSchemaStr = recordSchema.getField(DEFAULT_VALUE_FIELD_PROP).schema().toString(); + + TestUtils.assertCommand(parentControllerClient.createNewStore(storeName, "owner", keySchemaStr, valueSchemaStr)); + UpdateStoreQueryParams updateStoreParams1 = + new UpdateStoreQueryParams().setStorageQuotaInByte(Store.UNLIMITED_STORAGE_QUOTA).setPartitionCount(1); + TestUtils.assertCommand(parentControllerClient.updateStore(storeName, updateStoreParams1)); + + // Run a batch push first to create a batch version with 1 partition + try (VenicePushJob job = new VenicePushJob("Test push job batch with NR + A/A all fabrics", propsBatch)) { + job.run(); + } + + // wait until version is created and verify the partition count + TestUtils.waitForNonDeterministicAssertion(1, TimeUnit.MINUTES, () -> { + StoreResponse storeResponse = assertCommand(parentControllerClient.getStore(storeName)); + StoreInfo storeInfo = storeResponse.getStore(); + assertNotNull(storeInfo, "Store info is null."); + assertNull(storeInfo.getHybridStoreConfig(), "Hybrid store config is not null."); + assertNotNull(storeInfo.getVersion(1), "Version 1 is not present."); + Optional version = storeInfo.getVersion(1); + assertTrue(version.isPresent(), "Version 1 is not present."); + assertNull(version.get().getHybridStoreConfig(), "Version level hybrid store config is not null."); + assertEquals(version.get().getPartitionCount(), 1, "Partition count is not 1."); + }); + + // Update the store to have 3 partitions and convert it into a hybrid store + UpdateStoreQueryParams updateStoreParams = + new UpdateStoreQueryParams().setStorageQuotaInByte(Store.UNLIMITED_STORAGE_QUOTA) + .setPartitionCount(3) + .setHybridOffsetLagThreshold(TEST_TIMEOUT / 2) + .setHybridRewindSeconds(2L); + TestUtils.assertCommand(parentControllerClient.updateStore(storeName, updateStoreParams)); + + TestUtils.waitForNonDeterministicAssertion(1, TimeUnit.MINUTES, () -> { + StoreResponse storeResponse = assertCommand(parentControllerClient.getStore(storeName)); + StoreInfo storeInfo = storeResponse.getStore(); + assertNotNull(storeInfo, "Store info is null."); + assertNotNull(storeInfo.getHybridStoreConfig(), "Hybrid store config is null."); + // verify that there is just one version and it is batch version + assertEquals(storeInfo.getVersions().size(), 1, "Version count is not 1."); + Optional version = storeInfo.getVersion(1); + assertTrue(version.isPresent(), "Version 1 is not present."); + assertNull(version.get().getHybridStoreConfig(), "Version level hybrid store config is not null."); + assertEquals(version.get().getPartitionCount(), 1, "Partition count is not 1."); + }); + + // Push job step was disabled to reproduce the issue + // Run a full push to create a hybrid version with 3 partitions + try (VenicePushJob job = new VenicePushJob("push_job_to_create_hybrid_version", propsBatch)) { + job.run(); + } + + // wait until hybrid version is created and verify the partition count + TestUtils.waitForNonDeterministicAssertion(1, TimeUnit.MINUTES, () -> { + StoreResponse storeResponse = assertCommand(parentControllerClient.getStore(storeName)); + StoreInfo storeInfo = storeResponse.getStore(); + assertNotNull(storeInfo, "Store info is null."); + assertNotNull(storeInfo.getHybridStoreConfig(), "Hybrid store config is null."); + assertNotNull(storeInfo.getVersion(2), "Version 2 is not present."); + Optional version = storeInfo.getVersion(2); + assertTrue(version.isPresent(), "Version 2 is not present."); + assertNotNull(version.get().getHybridStoreConfig(), "Version level hybrid store config is null."); + assertEquals(version.get().getPartitionCount(), 3, "Partition count is not 3."); + }); + + VeniceSystemFactory factory = new VeniceSystemFactory(); + Map samzaConfig = IntegrationTestPushUtils.getSamzaProducerConfig(childDatacenters, 1, storeName); + VeniceSystemProducer veniceProducer = factory.getClosableProducer("venice", new MapConfig(samzaConfig), null); + veniceProducer.start(); + + PubSubTopicRepository pubSubTopicRepository = + childDatacenters.get(1).getClusters().get(clusterName).getPubSubTopicRepository(); + PubSubTopic realTimeTopic = pubSubTopicRepository.getTopic(Version.composeRealTimeTopic(storeName)); + + // wait for 120 secs and check producer getTopicName + TestUtils.waitForNonDeterministicAssertion(2, TimeUnit.MINUTES, () -> { + Assert.assertEquals(veniceProducer.getTopicName(), realTimeTopic.getName()); + }); + + try (TopicManager topicManager = + IntegrationTestPushUtils + .getTopicManagerRepo( + PUBSUB_OPERATION_TIMEOUT_MS_DEFAULT_VALUE, + 100, + 0l, + childDatacenters.get(1).getClusters().get(clusterName).getPubSubBrokerWrapper(), + pubSubTopicRepository) + .getLocalTopicManager()) { + int partitionCount = topicManager.getPartitionCount(realTimeTopic); + assertEquals(partitionCount, 3, "Partition count is not 3."); + } + } + } + /** * The purpose of this test is to verify that incremental push with RT policy succeeds when A/A is enabled in all * regions. And also incremental push can push to the closes kafka cluster from the grid using the SOURCE_GRID_CONFIG. diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestBatchForRocksDB.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestBatchForRocksDB.java index 7ab788b26d5..4015671b616 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestBatchForRocksDB.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestBatchForRocksDB.java @@ -1,6 +1,7 @@ package com.linkedin.venice.endToEnd; import static com.linkedin.davinci.store.rocksdb.RocksDBServerConfig.ROCKSDB_PLAIN_TABLE_FORMAT_ENABLED; +import static com.linkedin.venice.ConfigKeys.BLOB_TRANSFER_MANAGER_ENABLED; import static com.linkedin.venice.ConfigKeys.DATA_BASE_PATH; import static com.linkedin.venice.ConfigKeys.ENABLE_BLOB_TRANSFER; import static com.linkedin.venice.ConfigKeys.PERSISTENCE_TYPE; @@ -29,6 +30,7 @@ public VeniceClusterWrapper initializeVeniceCluster() { serverProperties.setProperty(SERVER_DATABASE_CHECKSUM_VERIFICATION_ENABLED, "true"); serverProperties.setProperty(SERVER_DATABASE_SYNC_BYTES_INTERNAL_FOR_DEFERRED_WRITE_MODE, "300"); serverProperties.setProperty(ENABLE_BLOB_TRANSFER, "true"); + serverProperties.setProperty(BLOB_TRANSFER_MANAGER_ENABLED, "true"); serverProperties.setProperty(DATA_BASE_PATH, BASE_DATA_PATH_1); veniceClusterWrapper.addVeniceServer(new Properties(), serverProperties); serverProperties.setProperty(DATA_BASE_PATH, BASE_DATA_PATH_2); diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestControllerGrpcEndpoints.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestControllerGrpcEndpoints.java new file mode 100644 index 00000000000..b1477109603 --- /dev/null +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestControllerGrpcEndpoints.java @@ -0,0 +1,100 @@ +package com.linkedin.venice.endToEnd; + +import static com.linkedin.venice.ConfigKeys.CONTROLLER_GRPC_SERVER_ENABLED; +import static com.linkedin.venice.integration.utils.VeniceClusterWrapper.DEFAULT_KEY_SCHEMA; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; + +import com.linkedin.venice.controllerapi.StoreResponse; +import com.linkedin.venice.integration.utils.ServiceFactory; +import com.linkedin.venice.integration.utils.VeniceClusterCreateOptions; +import com.linkedin.venice.integration.utils.VeniceClusterWrapper; +import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcRequest; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcResponse; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcRequest; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcResponse; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcRequest; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcResponse; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcServiceGrpc; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcServiceGrpc.VeniceControllerGrpcServiceBlockingStub; +import com.linkedin.venice.utils.TestUtils; +import com.linkedin.venice.utils.Utils; +import io.grpc.Grpc; +import io.grpc.InsecureChannelCredentials; +import io.grpc.ManagedChannel; +import java.util.Properties; +import org.testng.annotations.AfterClass; +import org.testng.annotations.BeforeClass; +import org.testng.annotations.Test; + + +public class TestControllerGrpcEndpoints { + private VeniceClusterWrapper veniceCluster; + + @BeforeClass(alwaysRun = true) + public void setUp() { + Properties properties = new Properties(); + properties.put(CONTROLLER_GRPC_SERVER_ENABLED, true); + VeniceClusterCreateOptions options = new VeniceClusterCreateOptions.Builder().numberOfControllers(1) + .numberOfRouters(1) + .numberOfServers(1) + .extraProperties(properties) + .build(); + veniceCluster = ServiceFactory.getVeniceCluster(options); + } + + @AfterClass(alwaysRun = true) + public void tearDown() { + Utils.closeQuietlyWithErrorLogged(veniceCluster); + } + + @Test + public void testGrpcEndpointsWithGrpcClient() { + String storeName = Utils.getUniqueString("test_grpc_store"); + String controllerGrpcUrl = veniceCluster.getLeaderVeniceController().getControllerGrpcUrl(); + ManagedChannel channel = Grpc.newChannelBuilder(controllerGrpcUrl, InsecureChannelCredentials.create()).build(); + VeniceControllerGrpcServiceBlockingStub blockingStub = VeniceControllerGrpcServiceGrpc.newBlockingStub(channel); + + // Test 1: getLeaderControllerDetails + LeaderControllerGrpcResponse grpcResponse = blockingStub.getLeaderController( + LeaderControllerGrpcRequest.newBuilder().setClusterName(veniceCluster.getClusterName()).build()); + assertEquals(grpcResponse.getHttpUrl(), veniceCluster.getLeaderVeniceController().getControllerUrl()); + assertEquals(grpcResponse.getGrpcUrl(), veniceCluster.getLeaderVeniceController().getControllerGrpcUrl()); + assertEquals( + grpcResponse.getSecureGrpcUrl(), + veniceCluster.getLeaderVeniceController().getControllerSecureGrpcUrl()); + + // Test 2: createStore + CreateStoreGrpcRequest createStoreGrpcRequest = CreateStoreGrpcRequest.newBuilder() + .setClusterStoreInfo( + ClusterStoreGrpcInfo.newBuilder() + .setClusterName(veniceCluster.getClusterName()) + .setStoreName(storeName) + .build()) + .setOwner("owner") + .setKeySchema(DEFAULT_KEY_SCHEMA) + .setValueSchema("\"string\"") + .build(); + + CreateStoreGrpcResponse response = blockingStub.createStore(createStoreGrpcRequest); + assertNotNull(response, "Response should not be null"); + assertNotNull(response.getClusterStoreInfo(), "ClusterStoreInfo should not be null"); + assertEquals(response.getClusterStoreInfo().getClusterName(), veniceCluster.getClusterName()); + assertEquals(response.getClusterStoreInfo().getStoreName(), storeName); + + veniceCluster.useControllerClient(controllerClient -> { + StoreResponse storeResponse = TestUtils.assertCommand(controllerClient.getStore(storeName)); + assertNotNull(storeResponse.getStore(), "Store should not be null"); + }); + + // Test 3: discover cluster + DiscoverClusterGrpcRequest discoverClusterGrpcRequest = + DiscoverClusterGrpcRequest.newBuilder().setStoreName(storeName).build(); + DiscoverClusterGrpcResponse discoverClusterGrpcResponse = + blockingStub.discoverClusterForStore(discoverClusterGrpcRequest); + assertNotNull(discoverClusterGrpcResponse, "Response should not be null"); + assertEquals(discoverClusterGrpcResponse.getStoreName(), storeName); + assertEquals(discoverClusterGrpcResponse.getClusterName(), veniceCluster.getClusterName()); + } +} diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestIntToStringRecordTransformer.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestIntToStringRecordTransformer.java index bfcb13a04d6..d0d9f58def7 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestIntToStringRecordTransformer.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestIntToStringRecordTransformer.java @@ -10,18 +10,13 @@ * Transforms int values to strings */ public class TestIntToStringRecordTransformer extends DaVinciRecordTransformer { - public TestIntToStringRecordTransformer(int storeVersion, boolean storeRecordsInDaVinci) { - super(storeVersion, storeRecordsInDaVinci); - } - - @Override - public Schema getKeySchema() { - return Schema.create(Schema.Type.INT); - } - - @Override - public Schema getOutputValueSchema() { - return Schema.create(Schema.Type.STRING); + public TestIntToStringRecordTransformer( + int storeVersion, + Schema keySchema, + Schema inputValueSchema, + Schema outputValueSchema, + boolean storeRecordsInDaVinci) { + super(storeVersion, keySchema, inputValueSchema, outputValueSchema, storeRecordsInDaVinci); } @Override diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestSkipResultRecordTransformer.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestSkipResultRecordTransformer.java index f787ce0d2ba..25a7fb39fca 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestSkipResultRecordTransformer.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestSkipResultRecordTransformer.java @@ -12,18 +12,13 @@ public class TestSkipResultRecordTransformer extends DaVinciRecordTransformer { private final Map inMemoryDB = new HashMap<>(); - public TestSkipResultRecordTransformer(int storeVersion, boolean storeRecordsInDaVinci) { - super(storeVersion, storeRecordsInDaVinci); - } - - @Override - public Schema getKeySchema() { - return Schema.create(Schema.Type.INT); - } - - @Override - public Schema getOutputValueSchema() { - return Schema.create(Schema.Type.STRING); + public TestSkipResultRecordTransformer( + int storeVersion, + Schema keySchema, + Schema inputValueSchema, + Schema outputValueSchema, + boolean storeRecordsInDaVinci) { + super(storeVersion, keySchema, inputValueSchema, outputValueSchema, storeRecordsInDaVinci); } @Override diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestStringRecordTransformer.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestStringRecordTransformer.java index ba13b4aa49b..aba2c0baf69 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestStringRecordTransformer.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestStringRecordTransformer.java @@ -12,18 +12,13 @@ public class TestStringRecordTransformer extends DaVinciRecordTransformer { private final Map inMemoryDB = new HashMap<>(); - public TestStringRecordTransformer(int storeVersion, boolean storeRecordsInDaVinci) { - super(storeVersion, storeRecordsInDaVinci); - } - - @Override - public Schema getKeySchema() { - return Schema.create(Schema.Type.INT); - } - - @Override - public Schema getOutputValueSchema() { - return Schema.create(Schema.Type.STRING); + public TestStringRecordTransformer( + int storeVersion, + Schema keySchema, + Schema inputValueSchema, + Schema outputValueSchema, + boolean storeRecordsInDaVinci) { + super(storeVersion, keySchema, inputValueSchema, outputValueSchema, storeRecordsInDaVinci); } @Override diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestUnchangedResultRecordTransformer.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestUnchangedResultRecordTransformer.java index 0f46fb12298..c1a5320fa08 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestUnchangedResultRecordTransformer.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/endToEnd/TestUnchangedResultRecordTransformer.java @@ -7,18 +7,14 @@ public class TestUnchangedResultRecordTransformer extends DaVinciRecordTransformer { - public TestUnchangedResultRecordTransformer(int storeVersion, boolean storeRecordsInDaVinci) { - super(storeVersion, storeRecordsInDaVinci); - } - - @Override - public Schema getKeySchema() { - return Schema.create(Schema.Type.INT); - } - - @Override - public Schema getOutputValueSchema() { - return Schema.create(Schema.Type.STRING); + public TestUnchangedResultRecordTransformer( + int storeVersion, + + Schema keySchema, + Schema inputValueSchema, + Schema outputValueSchema, + boolean storeRecordsInDaVinci) { + super(storeVersion, keySchema, inputValueSchema, outputValueSchema, storeRecordsInDaVinci); } @Override diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/ServiceFactory.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/ServiceFactory.java index e412b1af2bb..6a51152ef6e 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/ServiceFactory.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/ServiceFactory.java @@ -14,6 +14,7 @@ import com.linkedin.venice.client.store.ClientConfig; import com.linkedin.venice.controller.Admin; import com.linkedin.venice.controller.server.AdminSparkServer; +import com.linkedin.venice.controller.server.VeniceControllerRequestHandler; import com.linkedin.venice.controllerapi.ControllerRoute; import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.pubsub.PubSubClientsFactory; @@ -152,7 +153,8 @@ public static VeniceControllerWrapper getVeniceController(VeniceControllerCreate public static AdminSparkServer getMockAdminSparkServer( Admin admin, String cluster, - List bannedRoutes) { + List bannedRoutes, + VeniceControllerRequestHandler requestHandler) { return getService("MockAdminSparkServer", (serviceName) -> { Set clusters = new HashSet<>(); clusters.add(cluster); @@ -168,7 +170,8 @@ public static AdminSparkServer getMockAdminSparkServer( bannedRoutes, null, false, - new PubSubTopicRepository()); // Change this. + new PubSubTopicRepository(), + requestHandler); server.start(); return server; }); diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/VeniceClusterWrapper.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/VeniceClusterWrapper.java index 8692a77799c..712b252162e 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/VeniceClusterWrapper.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/VeniceClusterWrapper.java @@ -493,6 +493,10 @@ public synchronized List getVeniceServers() { return new ArrayList<>(veniceServerWrappers.values()); } + public synchronized VeniceServerWrapper getVeniceServerByPort(int port) { + return veniceServerWrappers.get(port); + } + public synchronized Map getNettyServerToGrpcAddress() { return nettyServerToGrpcAddress; } @@ -536,6 +540,21 @@ public final synchronized String getAllControllersURLs() { .collect(Collectors.joining(",")); } + /** + * Retrieves the gRPC URLs of all available Venice controllers as a comma-separated string. + * + * @return A comma-separated string of gRPC URLs for all controllers. If no controllers are available, + * the {@code externalControllerDiscoveryURL} is returned. + */ + public final synchronized String getAllControllersGrpcURLs() { + return veniceControllerWrappers.isEmpty() + ? externalControllerDiscoveryURL + : veniceControllerWrappers.values() + .stream() + .map(VeniceControllerWrapper::getControllerGrpcUrl) + .collect(Collectors.joining(",")); + } + public VeniceControllerWrapper getLeaderVeniceController() { return getLeaderVeniceController(60 * Time.MS_PER_SECOND); } diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/VeniceControllerWrapper.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/VeniceControllerWrapper.java index 78e51df38b7..61043bb205f 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/VeniceControllerWrapper.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/VeniceControllerWrapper.java @@ -15,6 +15,8 @@ import static com.linkedin.venice.ConfigKeys.CLUSTER_TO_SERVER_D2; import static com.linkedin.venice.ConfigKeys.CONCURRENT_INIT_ROUTINES_ENABLED; import static com.linkedin.venice.ConfigKeys.CONTROLLER_ADD_VERSION_VIA_ADMIN_PROTOCOL; +import static com.linkedin.venice.ConfigKeys.CONTROLLER_ADMIN_GRPC_PORT; +import static com.linkedin.venice.ConfigKeys.CONTROLLER_ADMIN_SECURE_GRPC_PORT; import static com.linkedin.venice.ConfigKeys.CONTROLLER_NAME; import static com.linkedin.venice.ConfigKeys.CONTROLLER_PARENT_MODE; import static com.linkedin.venice.ConfigKeys.CONTROLLER_SSL_ENABLED; @@ -109,6 +111,8 @@ public class VeniceControllerWrapper extends ProcessWrapper { private final boolean isParent; private final int port; private final int securePort; + private final int adminGrpcPort; + private final int adminSecureGrpcPort; private final String zkAddress; private final List d2ServerList; private final MetricsRepository metricsRepository; @@ -121,6 +125,8 @@ private VeniceControllerWrapper( VeniceController service, int port, int securePort, + int adminGrpcPort, + int adminSecureGrpcPort, List configs, boolean isParent, List d2ServerList, @@ -132,6 +138,8 @@ private VeniceControllerWrapper( this.isParent = isParent; this.port = port; this.securePort = securePort; + this.adminGrpcPort = adminGrpcPort; + this.adminSecureGrpcPort = adminSecureGrpcPort; this.zkAddress = zkAddress; this.d2ServerList = d2ServerList; this.metricsRepository = metricsRepository; @@ -142,6 +150,8 @@ static StatefulServiceProvider generateService(VeniceCo return (serviceName, dataDirectory) -> { int adminPort = TestUtils.getFreePort(); int adminSecurePort = TestUtils.getFreePort(); + int adminGrpcPort = TestUtils.getFreePort(); + int adminSecureGrpcPort = TestUtils.getFreePort(); List propertiesList = new ArrayList<>(); VeniceProperties extraProps = new VeniceProperties(options.getExtraProperties()); @@ -182,6 +192,8 @@ static StatefulServiceProvider generateService(VeniceCo .put(DEFAULT_REPLICA_FACTOR, options.getReplicationFactor()) .put(ADMIN_PORT, adminPort) .put(ADMIN_SECURE_PORT, adminSecurePort) + .put(CONTROLLER_ADMIN_GRPC_PORT, adminGrpcPort) + .put(CONTROLLER_ADMIN_SECURE_GRPC_PORT, adminSecureGrpcPort) .put(DEFAULT_PARTITION_SIZE, options.getPartitionSize()) .put(DEFAULT_NUMBER_OF_PARTITION, options.getNumberOfPartitions()) .put(DEFAULT_MAX_NUMBER_OF_PARTITIONS, options.getMaxNumberOfPartitions()) @@ -375,6 +387,8 @@ static StatefulServiceProvider generateService(VeniceCo veniceController, adminPort, adminSecurePort, + adminGrpcPort, + adminSecureGrpcPort, propertiesList, options.isParent(), d2ServerList, @@ -401,6 +415,22 @@ public int getSecurePort() { return securePort; } + public int getAdminGrpcPort() { + return adminGrpcPort; + } + + public int getAdminSecureGrpcPort() { + return adminSecureGrpcPort; + } + + public String getControllerGrpcUrl() { + return getHost() + ":" + getAdminGrpcPort(); + } + + public String getControllerSecureGrpcUrl() { + return getHost() + ":" + getAdminSecureGrpcPort(); + } + public String getControllerUrl() { return "http://" + getHost() + ":" + getPort(); } diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/ZkServerWrapper.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/ZkServerWrapper.java index 312c9295977..e309c247d13 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/ZkServerWrapper.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/integration/utils/ZkServerWrapper.java @@ -49,7 +49,7 @@ public class ZkServerWrapper extends ProcessWrapper { * The tick time can be low because this Zookeeper instance is intended to be used locally. */ private static final int TICK_TIME = 200; - private static final int MAX_SESSION_TIMEOUT = 10 * Time.MS_PER_SECOND; + private static final int MAX_SESSION_TIMEOUT = 60 * Time.MS_PER_SECOND; private static final int NUM_CONNECTIONS = 5000; private static final String CLIENT_PORT_PROP = "clientPort"; diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/router/TestRead.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/router/TestRead.java index 690c35a9ba8..79968b4f759 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/router/TestRead.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/router/TestRead.java @@ -156,6 +156,7 @@ public void setUp() throws VeniceClientException, ExecutionException, Interrupte extraProperties.put(ConfigKeys.ROUTER_HTTP2_INBOUND_ENABLED, isRouterHttp2Enabled()); extraProperties.put(ConfigKeys.SERVER_HTTP2_INBOUND_ENABLED, true); extraProperties.put(ConfigKeys.ROUTER_PER_STORE_ROUTER_QUOTA_BUFFER, 0.0); + extraProperties.put(ConfigKeys.PARTICIPANT_MESSAGE_STORE_ENABLED, false); veniceCluster = ServiceFactory.getVeniceCluster(1, 1, 1, 2, 100, true, false, extraProperties); routerAddr = veniceCluster.getRandomRouterSslURL(); diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/router/TestRetryQuotaRejection.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/router/TestRetryQuotaRejection.java index c3263451d1d..43f470f1364 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/router/TestRetryQuotaRejection.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/router/TestRetryQuotaRejection.java @@ -1,7 +1,6 @@ package com.linkedin.venice.router; -import static com.linkedin.venice.ConfigKeys.ROUTER_ENABLE_READ_THROTTLING; -import static com.linkedin.venice.ConfigKeys.SERVER_QUOTA_ENFORCEMENT_ENABLED; +import static com.linkedin.venice.ConfigKeys.*; import static org.testng.Assert.assertEquals; import com.linkedin.d2.balancer.D2Client; @@ -16,6 +15,7 @@ import com.linkedin.venice.helix.HelixReadOnlySchemaRepository; import com.linkedin.venice.integration.utils.D2TestUtils; import com.linkedin.venice.integration.utils.ServiceFactory; +import com.linkedin.venice.integration.utils.VeniceClusterCreateOptions; import com.linkedin.venice.integration.utils.VeniceClusterWrapper; import com.linkedin.venice.integration.utils.VeniceRouterWrapper; import com.linkedin.venice.integration.utils.VeniceServerWrapper; @@ -83,8 +83,17 @@ public void setUp() throws VeniceClientException, ExecutionException, Interrupte extraProperties.put(ConfigKeys.ROUTER_STORAGE_NODE_CLIENT_TYPE, StorageNodeClientType.APACHE_HTTP_ASYNC_CLIENT); extraProperties.put(ROUTER_ENABLE_READ_THROTTLING, false); extraProperties.put(SERVER_QUOTA_ENFORCEMENT_ENABLED, "true"); - - veniceCluster = ServiceFactory.getVeniceCluster(1, 1, 2, 2, 100, true, false, extraProperties); + extraProperties.put(PARTICIPANT_MESSAGE_STORE_ENABLED, "false"); + + veniceCluster = ServiceFactory.getVeniceCluster( + new VeniceClusterCreateOptions.Builder().numberOfControllers(1) + .numberOfServers(1) + .numberOfRouters(2) + .replicationFactor(2) + .partitionSize(100) + .sslToStorageNodes(true) + .extraProperties(extraProperties) + .build()); Properties serverProperties = new Properties(); Properties serverFeatureProperties = new Properties(); serverFeatureProperties.put(VeniceServerWrapper.SERVER_ENABLE_SSL, "true"); diff --git a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/server/VeniceServerTest.java b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/server/VeniceServerTest.java index 68aa27f4960..6fb1c6a615c 100644 --- a/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/server/VeniceServerTest.java +++ b/internal/venice-test-common/src/integrationTest/java/com/linkedin/venice/server/VeniceServerTest.java @@ -118,15 +118,13 @@ public void testCheckBeforeJoinCluster() { try (VeniceClusterWrapper cluster = ServiceFactory.getVeniceCluster(1, 1, 0)) { VeniceServerWrapper server = cluster.getVeniceServers().get(0); StorageEngineRepository repository = server.getVeniceServer().getStorageService().getStorageEngineRepository(); - Assert - .assertTrue(repository.getAllLocalStorageEngines().isEmpty(), "New node should not have any storage engine."); // Create a storage engine. String storeName = cluster.createStore(10); - Assert.assertNotEquals( - repository.getAllLocalStorageEngines().size(), - 0, - "We have created one storage engine for store: " + storeName); + String storeVersionName = Version.composeKafkaTopic(storeName, 1); + Assert.assertNotNull( + repository.getLocalStorageEngine(storeVersionName), + "Storage engine should be created for: " + storeVersionName); // Restart server, as server's info leave in Helix cluster, so we expect that all local storage would NOT be // deleted @@ -136,6 +134,9 @@ public void testCheckBeforeJoinCluster() { repository = server.getVeniceServer().getStorageService().getStorageEngineRepository(); Assert .assertNotEquals(repository.getAllLocalStorageEngines().size(), 0, "We should not cleanup the local storage"); + Assert.assertNotNull( + repository.getLocalStorageEngine(storeVersionName), + "Storage engine should be created for: " + storeVersionName); // Stop server, remove it from the cluster then restart. We expect that all local storage would be deleted. Once // the server join again. @@ -163,12 +164,11 @@ public void testCheckBeforeJoinCluster() { TestUtils.waitForNonDeterministicAssertion( 30, TimeUnit.SECONDS, - () -> Assert.assertTrue( + () -> Assert.assertNull( server.getVeniceServer() .getStorageService() .getStorageEngineRepository() - .getAllLocalStorageEngines() - .isEmpty(), + .getLocalStorageEngine(storeVersionName), "After removing the node from cluster, local storage should be cleaned up once the server join the cluster again.")); } } @@ -210,12 +210,12 @@ public void testStartServerAndShutdownWithPartitionAssignmentVerification() { Assert.assertTrue(server.getVeniceServer().isStarted()); StorageService storageService = server.getVeniceServer().getStorageService(); StorageEngineRepository repository = storageService.getStorageEngineRepository(); - Assert - .assertTrue(repository.getAllLocalStorageEngines().isEmpty(), "New node should not have any storage engine."); // Create a storage engine. String storeName = Version.composeKafkaTopic(cluster.createStore(1), 1); - Assert.assertEquals(repository.getAllLocalStorageEngines().size(), 1); + Assert.assertNotNull( + repository.getLocalStorageEngine(storeName), + "Storage engine should be created for: " + storeName); Assert.assertTrue(server.getVeniceServer().getHelixParticipationService().isRunning()); Assert.assertEquals(storageService.getStorageEngine(storeName).getPartitionIds().size(), 3); @@ -404,18 +404,14 @@ public void testDropStorePartitionAsynchronously() { StorageService storageService = server.getVeniceServer().getStorageService(); Assert.assertTrue(server.getVeniceServer().isStarted()); final StorageEngineRepository repository = storageService.getStorageEngineRepository(); - Assert - .assertTrue(repository.getAllLocalStorageEngines().isEmpty(), "New node should not have any storage engine."); - // Create a new store String storeName = cluster.createStore(1); String storeVersionName = Version.composeKafkaTopic(storeName, 1); - Assert.assertEquals(repository.getAllLocalStorageEngines().size(), 1); + Assert.assertNotNull( + storageService.getStorageEngine(storeVersionName), + "Storage engine should be created for: " + storeVersionName); Assert.assertTrue(server.getVeniceServer().getHelixParticipationService().isRunning()); Assert.assertEquals(storageService.getStorageEngine(storeVersionName).getPartitionIds().size(), 3); - - Assert.assertEquals(repository.getAllLocalStorageEngines().size(), 1); - String helixInstanceName = Utils.getHelixNodeIdentifier(Utils.getHostName(), server.getPort()); String instanceOperationReason = "Disable instance to remove all partitions assigned to it"; cluster.getLeaderVeniceController() @@ -429,6 +425,9 @@ public void testDropStorePartitionAsynchronously() { TestUtils.waitForNonDeterministicAssertion(10, TimeUnit.SECONDS, () -> { // All partitions should have been dropped asynchronously due to instance being disabled + Assert.assertNull( + storageService.getStorageEngine(storeVersionName), + "Storage engine: " + storeVersionName + " should have been dropped"); Assert.assertEquals(repository.getAllLocalStorageEngines().size(), 0); }); } @@ -445,14 +444,13 @@ public void testDropStorePartitionSynchronously() { StorageService storageService = server.getVeniceServer().getStorageService(); Assert.assertTrue(server.getVeniceServer().isStarted()); - final StorageEngineRepository repository = storageService.getStorageEngineRepository(); - Assert - .assertTrue(repository.getAllLocalStorageEngines().isEmpty(), "New node should not have any storage engine."); // Create a new store String storeName = cluster.createStore(1); String storeVersionName = Version.composeKafkaTopic(storeName, 1); - Assert.assertEquals(repository.getAllLocalStorageEngines().size(), 1); + Assert.assertNotNull( + storageService.getStorageEngine(storeVersionName), + "Storage engine should be created for: " + storeVersionName); Assert.assertTrue(server.getVeniceServer().getHelixParticipationService().isRunning()); Assert.assertEquals(storageService.getStorageEngine(storeVersionName).getPartitionIds().size(), 3); @@ -462,7 +460,9 @@ public void testDropStorePartitionSynchronously() { TestUtils.waitForNonDeterministicAssertion(10, TimeUnit.SECONDS, () -> { // All partitions should have been dropped synchronously - Assert.assertEquals(repository.getAllLocalStorageEngines().size(), 0); + Assert.assertNull( + storageService.getStorageEngine(storeVersionName), + "Storage engine: " + storeVersionName + " should have been dropped"); }); } } diff --git a/internal/venice-test-common/src/jmh/java/com/linkedin/venice/benchmark/TokenBucketBenchmark.java b/internal/venice-test-common/src/jmh/java/com/linkedin/venice/benchmark/TokenBucketBenchmark.java index 51b0cd241a8..c11d0832a5e 100644 --- a/internal/venice-test-common/src/jmh/java/com/linkedin/venice/benchmark/TokenBucketBenchmark.java +++ b/internal/venice-test-common/src/jmh/java/com/linkedin/venice/benchmark/TokenBucketBenchmark.java @@ -1,6 +1,6 @@ package com.linkedin.venice.benchmark; -import static java.util.concurrent.TimeUnit.*; +import static java.util.concurrent.TimeUnit.SECONDS; import com.linkedin.venice.throttle.TokenBucket; import java.text.DecimalFormat; diff --git a/internal/venice-test-common/src/main/java/com/linkedin/venice/utils/TestUtils.java b/internal/venice-test-common/src/main/java/com/linkedin/venice/utils/TestUtils.java index b135c57ad8e..1f06142bf9f 100644 --- a/internal/venice-test-common/src/main/java/com/linkedin/venice/utils/TestUtils.java +++ b/internal/venice-test-common/src/main/java/com/linkedin/venice/utils/TestUtils.java @@ -639,6 +639,8 @@ public static Properties getPropertiesForControllerConfig() { properties.put(ConfigKeys.DEFAULT_NUMBER_OF_PARTITION, "1"); properties.put(ConfigKeys.ADMIN_PORT, TestUtils.getFreePort()); properties.put(ConfigKeys.ADMIN_SECURE_PORT, TestUtils.getFreePort()); + properties.put(ConfigKeys.CONTROLLER_ADMIN_GRPC_PORT, TestUtils.getFreePort()); + properties.put(ConfigKeys.CONTROLLER_ADMIN_SECURE_GRPC_PORT, TestUtils.getFreePort()); return properties; } diff --git a/internal/venice-test-common/src/main/java/com/linkedin/venice/utils/TestWriteUtils.java b/internal/venice-test-common/src/main/java/com/linkedin/venice/utils/TestWriteUtils.java index 71f63834567..3b687c904e9 100644 --- a/internal/venice-test-common/src/main/java/com/linkedin/venice/utils/TestWriteUtils.java +++ b/internal/venice-test-common/src/main/java/com/linkedin/venice/utils/TestWriteUtils.java @@ -48,6 +48,7 @@ import java.util.Objects; import java.util.Properties; import java.util.function.Consumer; +import java.util.function.Function; import org.apache.avro.Schema; import org.apache.avro.file.DataFileWriter; import org.apache.avro.generic.GenericData; @@ -80,6 +81,10 @@ public class TestWriteUtils { AvroCompatibilityHelper.parse(loadSchemaFileFromResource("valueSchema/User.avsc")); public static final Schema USER_WITH_DEFAULT_SCHEMA = AvroCompatibilityHelper.parse(loadSchemaFileFromResource("valueSchema/UserWithDefault.avsc")); + + public static final Schema SINGLE_FIELD_RECORD_SCHEMA = + AvroCompatibilityHelper.parse(loadSchemaFileFromResource("valueSchema/SingleFieldRecord.avsc")); + public static final Schema SIMPLE_USER_WITH_DEFAULT_SCHEMA = AvroCompatibilityHelper.parse(loadSchemaFileFromResource("valueSchema/SimpleUserWithDefault.avsc")); public static final Schema USER_WITH_FLOAT_ARRAY_SCHEMA = @@ -353,17 +358,27 @@ public static Schema writeEmptyAvroFile(File parentDir, String fileName, Schema } public static Schema writeSimpleAvroFileWithStringToNameRecordV1Schema(File parentDir) throws IOException { - return writeAvroFile(parentDir, "string2record.avro", STRING_TO_NAME_RECORD_V1_SCHEMA, (recordSchema, writer) -> { - String firstName = "first_name_"; - String lastName = "last_name_"; + String firstName = "first_name_"; + String lastName = "last_name_"; + + return writeSimpleAvroFile(parentDir, STRING_TO_NAME_RECORD_V1_SCHEMA, i -> { + GenericRecord keyValueRecord = new GenericData.Record(STRING_TO_NAME_RECORD_V1_SCHEMA); + keyValueRecord.put(DEFAULT_KEY_FIELD_PROP, String.valueOf(i)); // Key + GenericRecord valueRecord = new GenericData.Record(NAME_RECORD_V1_SCHEMA); + valueRecord.put("firstName", firstName + i); + valueRecord.put("lastName", lastName + i); + keyValueRecord.put(DEFAULT_VALUE_FIELD_PROP, valueRecord); // Value + return keyValueRecord; + }); + } + + public static Schema writeSimpleAvroFile( + File parentDir, + Schema schema, + Function recordProvider) throws IOException { + return writeAvroFile(parentDir, "string2record.avro", schema, (recordSchema, writer) -> { for (int i = 1; i <= DEFAULT_USER_DATA_RECORD_COUNT; ++i) { - GenericRecord keyValueRecord = new GenericData.Record(recordSchema); - keyValueRecord.put(DEFAULT_KEY_FIELD_PROP, String.valueOf(i)); // Key - GenericRecord valueRecord = new GenericData.Record(NAME_RECORD_V1_SCHEMA); - valueRecord.put("firstName", firstName + i); - valueRecord.put("lastName", lastName + i); - keyValueRecord.put(DEFAULT_VALUE_FIELD_PROP, valueRecord); // Value - writer.append(keyValueRecord); + writer.append(recordProvider.apply(i)); } }); } diff --git a/internal/venice-test-common/src/main/resources/valueSchema/SingleFieldRecord.avsc b/internal/venice-test-common/src/main/resources/valueSchema/SingleFieldRecord.avsc new file mode 100644 index 00000000000..8757fa694fb --- /dev/null +++ b/internal/venice-test-common/src/main/resources/valueSchema/SingleFieldRecord.avsc @@ -0,0 +1,10 @@ +{ + "type" : "record", + "name" : "SingleFieldRecord", + "namespace" : "example.avro", + "fields" : [ { + "name" : "key", + "type" : "string", + "default" : "" + } ] +} diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/Admin.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/Admin.java index 2358c09762b..9ea86030a1f 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/Admin.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/Admin.java @@ -9,7 +9,6 @@ import com.linkedin.venice.controllerapi.UpdateClusterConfigQueryParams; import com.linkedin.venice.controllerapi.UpdateStoragePersonaQueryParams; import com.linkedin.venice.controllerapi.UpdateStoreQueryParams; -import com.linkedin.venice.exceptions.VeniceNoStoreException; import com.linkedin.venice.helix.HelixReadOnlyStoreConfigRepository; import com.linkedin.venice.helix.HelixReadOnlyZKSharedSchemaRepository; import com.linkedin.venice.helix.HelixReadOnlyZKSharedSystemStoreRepository; @@ -289,26 +288,9 @@ Version incrementVersionIdempotent( String targetedRegions, int repushSourceVersion); - String getRealTimeTopic(String clusterName, Store store); + Version getIncrementalPushVersion(String clusterName, String storeName, String pushJobId); - default String getRealTimeTopic(String clusterName, String storeName) { - Store store = getStore(clusterName, storeName); - if (store == null) { - throw new VeniceNoStoreException(storeName, clusterName); - } - return getRealTimeTopic(clusterName, store); - } - - String getSeparateRealTimeTopic(String clusterName, String storeName); - - /** - * Right now, it will return the latest version recorded in parent controller. There are a couple of edge cases. - * 1. If a push fails in some colos, the version will be inconsistent among colos - * 2. If rollback happens, latest version will not be the current version. - * - * TODO: figure out how we'd like to cover these edge cases - */ - Version getIncrementalPushVersion(String clusterName, String storeName); + Version getReferenceVersionForStreamingWrites(String clusterName, String storeName, String pushJobId); int getCurrentVersion(String clusterName, String storeName); diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/ControllerRequestHandlerDependencies.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/ControllerRequestHandlerDependencies.java new file mode 100644 index 00000000000..fec9a339285 --- /dev/null +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/ControllerRequestHandlerDependencies.java @@ -0,0 +1,208 @@ +package com.linkedin.venice.controller; + +import com.linkedin.venice.SSLConfig; +import com.linkedin.venice.acl.DynamicAccessController; +import com.linkedin.venice.acl.NoOpDynamicAccessController; +import com.linkedin.venice.controller.server.VeniceControllerAccessManager; +import com.linkedin.venice.controllerapi.ControllerRoute; +import com.linkedin.venice.pubsub.PubSubTopicRepository; +import com.linkedin.venice.utils.VeniceProperties; +import io.tehuti.metrics.MetricsRepository; +import java.util.Collections; +import java.util.List; +import java.util.Set; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * Dependencies for VeniceControllerRequestHandler + */ +public class ControllerRequestHandlerDependencies { + private static final Logger LOGGER = LogManager.getLogger(ControllerRequestHandlerDependencies.class); + private final Admin admin; + private final boolean enforceSSL; + private final boolean sslEnabled; + private final boolean checkReadMethodForKafka; + private final SSLConfig sslConfig; + private final DynamicAccessController accessController; + private final List disabledRoutes; + private final Set clusters; + private final boolean disableParentRequestTopicForStreamPushes; + private final PubSubTopicRepository pubSubTopicRepository; + private final MetricsRepository metricsRepository; + private final VeniceProperties veniceProperties; + private final VeniceControllerAccessManager controllerAccessManager; + + private ControllerRequestHandlerDependencies(Builder builder) { + this.admin = builder.admin; + this.enforceSSL = builder.enforceSSL; + this.sslEnabled = builder.sslEnabled; + this.checkReadMethodForKafka = builder.checkReadMethodForKafka; + this.sslConfig = builder.sslConfig; + this.accessController = builder.accessController; + this.disabledRoutes = builder.disabledRoutes; + this.clusters = builder.clusters; + this.disableParentRequestTopicForStreamPushes = builder.disableParentRequestTopicForStreamPushes; + this.pubSubTopicRepository = builder.pubSubTopicRepository; + this.metricsRepository = builder.metricsRepository; + this.veniceProperties = builder.veniceProperties; + this.controllerAccessManager = builder.controllerAccessManager; + } + + public Admin getAdmin() { + return admin; + } + + public Set getClusters() { + return clusters; + } + + public boolean isEnforceSSL() { + return enforceSSL; + } + + public boolean isSslEnabled() { + return sslEnabled; + } + + public boolean isCheckReadMethodForKafka() { + return checkReadMethodForKafka; + } + + public SSLConfig getSslConfig() { + return sslConfig; + } + + public DynamicAccessController getAccessController() { + return accessController; + } + + public List getDisabledRoutes() { + return disabledRoutes; + } + + public boolean isDisableParentRequestTopicForStreamPushes() { + return disableParentRequestTopicForStreamPushes; + } + + public PubSubTopicRepository getPubSubTopicRepository() { + return pubSubTopicRepository; + } + + public MetricsRepository getMetricsRepository() { + return metricsRepository; + } + + public VeniceProperties getVeniceProperties() { + return veniceProperties; + } + + public VeniceControllerAccessManager getControllerAccessManager() { + return controllerAccessManager; + } + + // Builder class for VeniceControllerRequestHandlerDependencies + public static class Builder { + private Admin admin; + private boolean enforceSSL; + private boolean sslEnabled; + private boolean checkReadMethodForKafka; + private SSLConfig sslConfig; + private DynamicAccessController accessController; + private List disabledRoutes; + private Set clusters; + private boolean disableParentRequestTopicForStreamPushes; + private PubSubTopicRepository pubSubTopicRepository; + private MetricsRepository metricsRepository; + private VeniceProperties veniceProperties; + private VeniceControllerAccessManager controllerAccessManager; + + public Builder setAdmin(Admin admin) { + this.admin = admin; + return this; + } + + public Builder setClusters(Set clusters) { + this.clusters = clusters; + return this; + } + + public Builder setEnforceSSL(boolean enforceSSL) { + this.enforceSSL = enforceSSL; + return this; + } + + public Builder setSslEnabled(boolean sslEnabled) { + this.sslEnabled = sslEnabled; + return this; + } + + public Builder setCheckReadMethodForKafka(boolean checkReadMethodForKafka) { + this.checkReadMethodForKafka = checkReadMethodForKafka; + return this; + } + + public Builder setSslConfig(SSLConfig sslConfig) { + this.sslConfig = sslConfig; + return this; + } + + public Builder setAccessController(DynamicAccessController accessController) { + this.accessController = accessController; + return this; + } + + public Builder setDisabledRoutes(List disabledRoutes) { + this.disabledRoutes = disabledRoutes; + return this; + } + + public Builder setDisableParentRequestTopicForStreamPushes(boolean disableParentRequestTopicForStreamPushes) { + this.disableParentRequestTopicForStreamPushes = disableParentRequestTopicForStreamPushes; + return this; + } + + public Builder setPubSubTopicRepository(PubSubTopicRepository pubSubTopicRepository) { + this.pubSubTopicRepository = pubSubTopicRepository; + return this; + } + + public Builder setMetricsRepository(MetricsRepository metricsRepository) { + this.metricsRepository = metricsRepository; + return this; + } + + public Builder setVeniceProperties(VeniceProperties veniceProperties) { + this.veniceProperties = veniceProperties; + return this; + } + + private void verifyAndAddDefaults() { + if (admin == null) { + throw new IllegalArgumentException("admin is mandatory dependencies for VeniceControllerRequestHandler"); + } + if (pubSubTopicRepository == null) { + pubSubTopicRepository = new PubSubTopicRepository(); + } + if (disabledRoutes == null) { + disabledRoutes = Collections.emptyList(); + } + if (sslEnabled && sslConfig == null) { + throw new IllegalArgumentException("sslConfig is mandatory when sslEnabled is true"); + } + if (accessController == null || !sslEnabled) { + String reason = (accessController == null ? "access controller is not configured" : "") + + (accessController == null && !sslEnabled ? " and " : "") + (!sslEnabled ? "SSL is disabled" : ""); + LOGGER.info("Defaulting to NoOpDynamicAccessController because {}.", reason); + accessController = NoOpDynamicAccessController.INSTANCE; + } + controllerAccessManager = new VeniceControllerAccessManager(accessController); + } + + public ControllerRequestHandlerDependencies build() { + verifyAndAddDefaults(); + return new ControllerRequestHandlerDependencies(this); + } + } +} diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceController.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceController.java index 2d07dd400bc..7dabd376b82 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceController.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceController.java @@ -11,13 +11,20 @@ import com.linkedin.venice.controller.kafka.TopicCleanupService; import com.linkedin.venice.controller.kafka.TopicCleanupServiceForParentController; import com.linkedin.venice.controller.server.AdminSparkServer; +import com.linkedin.venice.controller.server.VeniceControllerGrpcServiceImpl; +import com.linkedin.venice.controller.server.VeniceControllerRequestHandler; +import com.linkedin.venice.controller.server.grpc.ControllerGrpcSslSessionInterceptor; +import com.linkedin.venice.controller.server.grpc.ParentControllerRegionValidationInterceptor; import com.linkedin.venice.controller.stats.TopicCleanupServiceStats; import com.linkedin.venice.controller.supersetschema.SupersetSchemaGenerator; import com.linkedin.venice.controller.systemstore.SystemStoreRepairService; import com.linkedin.venice.d2.D2ClientFactory; import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.grpc.VeniceGrpcServer; +import com.linkedin.venice.grpc.VeniceGrpcServerConfig; import com.linkedin.venice.pubsub.PubSubClientsFactory; import com.linkedin.venice.pubsub.PubSubTopicRepository; +import com.linkedin.venice.security.SSLFactory; import com.linkedin.venice.serialization.avro.AvroProtocolDefinition; import com.linkedin.venice.service.AbstractVeniceService; import com.linkedin.venice.service.ICProvider; @@ -25,13 +32,18 @@ import com.linkedin.venice.servicediscovery.ServiceDiscoveryAnnouncer; import com.linkedin.venice.system.store.ControllerClientBackedSystemSchemaInitializer; import com.linkedin.venice.utils.PropertyBuilder; +import com.linkedin.venice.utils.SslUtils; import com.linkedin.venice.utils.Utils; import com.linkedin.venice.utils.VeniceProperties; +import com.linkedin.venice.utils.concurrent.BlockingQueueType; +import com.linkedin.venice.utils.concurrent.ThreadPoolFactory; +import io.grpc.ServerInterceptor; import io.tehuti.metrics.MetricsRepository; import java.util.ArrayList; import java.util.Collections; import java.util.List; import java.util.Optional; +import java.util.concurrent.ThreadPoolExecutor; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -41,19 +53,27 @@ */ public class VeniceController { private static final Logger LOGGER = LogManager.getLogger(VeniceController.class); + private static final String CONTROLLER_GRPC_SERVER_THREAD_NAME = "ControllerGrpcServer"; + static final String CONTROLLER_SERVICE_NAME = "venice-controller"; // services - private VeniceControllerService controllerService; - private AdminSparkServer adminServer; - private AdminSparkServer secureAdminServer; - private TopicCleanupService topicCleanupService; - private Optional storeBackupVersionCleanupService; + private final VeniceControllerService controllerService; + private final AdminSparkServer adminServer; + private final AdminSparkServer secureAdminServer; + private VeniceGrpcServer adminGrpcServer; + private VeniceGrpcServer adminSecureGrpcServer; + private final TopicCleanupService topicCleanupService; + private final Optional storeBackupVersionCleanupService; + + private final Optional disabledPartitionEnablerService; + private final Optional unusedValueSchemaCleanupService; - private Optional disabledPartitionEnablerService; - private Optional unusedValueSchemaCleanupService; + private final Optional storeGraveyardCleanupService; + private final Optional systemStoreRepairService; - private Optional storeGraveyardCleanupService; - private Optional systemStoreRepairService; + private VeniceControllerRequestHandler secureRequestHandler; + private VeniceControllerRequestHandler unsecureRequestHandler; + private ThreadPoolExecutor grpcExecutor = null; private final boolean sslEnabled; private final VeniceControllerMultiClusterConfig multiClusterConfigs; @@ -66,9 +86,8 @@ public class VeniceController { private final Optional routerClientConfig; private final Optional icProvider; private final Optional externalSupersetSchemaGenerator; - private final PubSubTopicRepository pubSubTopicRepository = new PubSubTopicRepository(); + private final PubSubTopicRepository pubSubTopicRepository; private final PubSubClientsFactory pubSubClientsFactory; - static final String CONTROLLER_SERVICE_NAME = "venice-controller"; /** * Allocates a new {@code VeniceController} object. @@ -132,11 +151,26 @@ public VeniceController(VeniceControllerContext ctx) { long serviceDiscoveryRegistrationRetryMS = multiClusterConfigs.getServiceDiscoveryRegistrationRetryMS(); this.asyncRetryingServiceDiscoveryAnnouncer = new AsyncRetryingServiceDiscoveryAnnouncer(serviceDiscoveryAnnouncers, serviceDiscoveryRegistrationRetryMS); - createServices(); + this.pubSubTopicRepository = new PubSubTopicRepository(); + this.controllerService = createControllerService(); + this.adminServer = createAdminServer(false); + this.secureAdminServer = sslEnabled ? createAdminServer(true) : null; + this.topicCleanupService = createTopicCleanupService(); + this.storeBackupVersionCleanupService = createStoreBackupVersionCleanupService(); + this.disabledPartitionEnablerService = createDisabledPartitionEnablerService(); + this.unusedValueSchemaCleanupService = createUnusedValueSchemaCleanupService(); + this.storeGraveyardCleanupService = createStoreGraveyardCleanupService(); + this.systemStoreRepairService = createSystemStoreRepairService(); + if (multiClusterConfigs.isGrpcServerEnabled()) { + initializeGrpcServer(); + } + + // Run before enabling controller in helix so leadership won't hand back to this controller during schema requests. + initializeSystemSchema(controllerService.getVeniceHelixAdmin()); } - private void createServices() { - controllerService = new VeniceControllerService( + private VeniceControllerService createControllerService() { + VeniceControllerService veniceControllerService = new VeniceControllerService( multiClusterConfigs, metricsRepository, sslEnabled, @@ -149,88 +183,152 @@ private void createServices() { externalSupersetSchemaGenerator, pubSubTopicRepository, pubSubClientsFactory); + Admin admin = veniceControllerService.getVeniceHelixAdmin(); + if (multiClusterConfigs.isParent() && !(admin instanceof VeniceParentHelixAdmin)) { + throw new VeniceException( + "'VeniceParentHelixAdmin' is expected of the returned 'Admin' from 'VeniceControllerService#getVeniceHelixAdmin' in parent mode"); + } + unsecureRequestHandler = new VeniceControllerRequestHandler(buildRequestHandlerDependencies(admin, false)); + if (sslEnabled) { + secureRequestHandler = new VeniceControllerRequestHandler(buildRequestHandlerDependencies(admin, true)); + } + return veniceControllerService; + } - adminServer = new AdminSparkServer( - // no need to pass the hostname, we are binding to all the addresses - multiClusterConfigs.getAdminPort(), + AdminSparkServer createAdminServer(boolean secure) { + return new AdminSparkServer( + secure ? multiClusterConfigs.getAdminSecurePort() : multiClusterConfigs.getAdminPort(), controllerService.getVeniceHelixAdmin(), metricsRepository, multiClusterConfigs.getClusters(), - multiClusterConfigs.isControllerEnforceSSLOnly(), - Optional.empty(), - false, - Optional.empty(), + secure || multiClusterConfigs.isControllerEnforceSSLOnly(), + secure ? multiClusterConfigs.getSslConfig() : Optional.empty(), + secure && multiClusterConfigs.adminCheckReadMethodForKafka(), + accessController, multiClusterConfigs.getDisabledRoutes(), multiClusterConfigs.getCommonConfig().getJettyConfigOverrides(), - // TODO: Builder pattern or just pass the config object here? multiClusterConfigs.getCommonConfig().isDisableParentRequestTopicForStreamPushes(), - pubSubTopicRepository); - if (sslEnabled) { - /** - * SSL enabled AdminSparkServer uses a different port number than the regular service. - */ - secureAdminServer = new AdminSparkServer( - multiClusterConfigs.getAdminSecurePort(), - controllerService.getVeniceHelixAdmin(), - metricsRepository, - multiClusterConfigs.getClusters(), - true, - multiClusterConfigs.getSslConfig(), - multiClusterConfigs.adminCheckReadMethodForKafka(), - accessController, - multiClusterConfigs.getDisabledRoutes(), - multiClusterConfigs.getCommonConfig().getJettyConfigOverrides(), - multiClusterConfigs.getCommonConfig().isDisableParentRequestTopicForStreamPushes(), - pubSubTopicRepository); - } - storeBackupVersionCleanupService = Optional.empty(); - storeGraveyardCleanupService = Optional.empty(); - systemStoreRepairService = Optional.empty(); - disabledPartitionEnablerService = Optional.empty(); - unusedValueSchemaCleanupService = Optional.empty(); + pubSubTopicRepository, + secure ? secureRequestHandler : unsecureRequestHandler); + } + private TopicCleanupService createTopicCleanupService() { Admin admin = controllerService.getVeniceHelixAdmin(); if (multiClusterConfigs.isParent()) { - topicCleanupService = new TopicCleanupServiceForParentController( + return new TopicCleanupServiceForParentController( admin, multiClusterConfigs, pubSubTopicRepository, new TopicCleanupServiceStats(metricsRepository), pubSubClientsFactory); - if (!(admin instanceof VeniceParentHelixAdmin)) { - throw new VeniceException( - "'VeniceParentHelixAdmin' is expected of the returned 'Admin' from 'VeniceControllerService#getVeniceHelixAdmin' in parent mode"); - } - storeGraveyardCleanupService = - Optional.of(new StoreGraveyardCleanupService((VeniceParentHelixAdmin) admin, multiClusterConfigs)); - LOGGER.info("StoreGraveyardCleanupService is enabled"); - if (multiClusterConfigs.getCommonConfig().isParentSystemStoreRepairServiceEnabled()) { - systemStoreRepairService = Optional - .of(new SystemStoreRepairService((VeniceParentHelixAdmin) admin, multiClusterConfigs, metricsRepository)); - LOGGER.info("SystemStoreRepairServiceEnabled is enabled"); - } - this.unusedValueSchemaCleanupService = - Optional.of(new UnusedValueSchemaCleanupService(multiClusterConfigs, (VeniceParentHelixAdmin) admin)); } else { - topicCleanupService = new TopicCleanupService( + return new TopicCleanupService( admin, multiClusterConfigs, pubSubTopicRepository, new TopicCleanupServiceStats(metricsRepository), pubSubClientsFactory); - if (!(admin instanceof VeniceHelixAdmin)) { - throw new VeniceException( - "'VeniceHelixAdmin' is expected of the returned 'Admin' from 'VeniceControllerService#getVeniceHelixAdmin' in child mode"); - } - storeBackupVersionCleanupService = Optional + } + } + + private Optional createStoreBackupVersionCleanupService() { + if (!multiClusterConfigs.isParent()) { + Admin admin = controllerService.getVeniceHelixAdmin(); + return Optional .of(new StoreBackupVersionCleanupService((VeniceHelixAdmin) admin, multiClusterConfigs, metricsRepository)); - LOGGER.info("StoreBackupVersionCleanupService is enabled"); + } + return Optional.empty(); + } - disabledPartitionEnablerService = - Optional.of(new DisabledPartitionEnablerService((VeniceHelixAdmin) admin, multiClusterConfigs)); + private Optional createDisabledPartitionEnablerService() { + if (!multiClusterConfigs.isParent()) { + Admin admin = controllerService.getVeniceHelixAdmin(); + return Optional.of(new DisabledPartitionEnablerService((VeniceHelixAdmin) admin, multiClusterConfigs)); } - // Run before enabling controller in helix so leadership won't hand back to this controller during schema requests. - initializeSystemSchema(controllerService.getVeniceHelixAdmin()); + return Optional.empty(); + } + + private Optional createStoreGraveyardCleanupService() { + if (multiClusterConfigs.isParent()) { + Admin admin = controllerService.getVeniceHelixAdmin(); + return Optional.of(new StoreGraveyardCleanupService((VeniceParentHelixAdmin) admin, multiClusterConfigs)); + } + return Optional.empty(); + } + + private Optional createSystemStoreRepairService() { + if (multiClusterConfigs.isParent() + && multiClusterConfigs.getCommonConfig().isParentSystemStoreRepairServiceEnabled()) { + Admin admin = controllerService.getVeniceHelixAdmin(); + return Optional + .of(new SystemStoreRepairService((VeniceParentHelixAdmin) admin, multiClusterConfigs, metricsRepository)); + } + return Optional.empty(); + } + + private Optional createUnusedValueSchemaCleanupService() { + if (multiClusterConfigs.isParent()) { + Admin admin = controllerService.getVeniceHelixAdmin(); + return Optional.of(new UnusedValueSchemaCleanupService(multiClusterConfigs, (VeniceParentHelixAdmin) admin)); + } + return Optional.empty(); + } + + // package-private for testing + private void initializeGrpcServer() { + LOGGER.info("Initializing gRPC server as it is enabled for the controller..."); + ParentControllerRegionValidationInterceptor parentControllerRegionValidationInterceptor = + new ParentControllerRegionValidationInterceptor(controllerService.getVeniceHelixAdmin()); + List interceptors = new ArrayList<>(2); + interceptors.add(parentControllerRegionValidationInterceptor); + + VeniceControllerGrpcServiceImpl grpcService = new VeniceControllerGrpcServiceImpl(unsecureRequestHandler); + grpcExecutor = ThreadPoolFactory.createThreadPool( + multiClusterConfigs.getGrpcServerThreadCount(), + CONTROLLER_GRPC_SERVER_THREAD_NAME, + Integer.MAX_VALUE, + BlockingQueueType.LINKED_BLOCKING_QUEUE); + + adminGrpcServer = new VeniceGrpcServer( + new VeniceGrpcServerConfig.Builder().setPort(multiClusterConfigs.getAdminGrpcPort()) + .setService(grpcService) + .setExecutor(grpcExecutor) + .setInterceptors(interceptors) + .build()); + + if (sslEnabled) { + interceptors.add(new ControllerGrpcSslSessionInterceptor()); + SSLFactory sslFactory = SslUtils.getSSLFactory( + multiClusterConfigs.getSslConfig().get().getSslProperties(), + multiClusterConfigs.getSslFactoryClassName()); + VeniceControllerGrpcServiceImpl secureGrpcService = new VeniceControllerGrpcServiceImpl(secureRequestHandler); + adminSecureGrpcServer = new VeniceGrpcServer( + new VeniceGrpcServerConfig.Builder().setPort(multiClusterConfigs.getAdminSecureGrpcPort()) + .setService(secureGrpcService) + .setExecutor(grpcExecutor) + .setSslFactory(sslFactory) + .setInterceptors(interceptors) + .build()); + } + } + + // package-private for testing + ControllerRequestHandlerDependencies buildRequestHandlerDependencies(Admin admin, boolean secure) { + ControllerRequestHandlerDependencies.Builder builder = + new ControllerRequestHandlerDependencies.Builder().setAdmin(admin) + .setMetricsRepository(metricsRepository) + .setClusters(multiClusterConfigs.getClusters()) + .setDisabledRoutes(multiClusterConfigs.getDisabledRoutes()) + .setVeniceProperties(multiClusterConfigs.getCommonConfig().getJettyConfigOverrides()) + .setDisableParentRequestTopicForStreamPushes( + multiClusterConfigs.getCommonConfig().isDisableParentRequestTopicForStreamPushes()) + .setPubSubTopicRepository(pubSubTopicRepository) + .setSslConfig(secure ? multiClusterConfigs.getSslConfig().orElse(null) : null) + .setSslEnabled(secure) + .setCheckReadMethodForKafka(secure && multiClusterConfigs.adminCheckReadMethodForKafka()) + .setAccessController(secure ? accessController.orElse(null) : null) + .setEnforceSSL(secure || multiClusterConfigs.isControllerEnforceSSLOnly()); + return builder.build(); } /** @@ -255,6 +353,12 @@ public void start() { disabledPartitionEnablerService.ifPresent(AbstractVeniceService::start); // register with service discovery at the end asyncRetryingServiceDiscoveryAnnouncer.register(); + if (adminGrpcServer != null) { + adminGrpcServer.start(); + } + if (adminSecureGrpcServer != null) { + adminSecureGrpcServer.start(); + } LOGGER.info("Controller is started."); } @@ -311,6 +415,16 @@ public void stop() { unusedValueSchemaCleanupService.ifPresent(Utils::closeQuietlyWithErrorLogged); storeBackupVersionCleanupService.ifPresent(Utils::closeQuietlyWithErrorLogged); disabledPartitionEnablerService.ifPresent(Utils::closeQuietlyWithErrorLogged); + if (adminGrpcServer != null) { + adminGrpcServer.stop(); + } + if (adminSecureGrpcServer != null) { + adminSecureGrpcServer.stop(); + } + if (grpcExecutor != null) { + LOGGER.info("Shutting down gRPC executor"); + grpcExecutor.shutdown(); + } Utils.closeQuietlyWithErrorLogged(topicCleanupService); Utils.closeQuietlyWithErrorLogged(secureAdminServer); Utils.closeQuietlyWithErrorLogged(adminServer); @@ -371,4 +485,41 @@ private static void addShutdownHook(VeniceController controller, String zkAddres D2ClientFactory.release(zkAddress); })); } + + // helper method to aid in testing + AdminSparkServer getSecureAdminServer() { + return secureAdminServer; + } + + VeniceGrpcServer getAdminSecureGrpcServer() { + return adminSecureGrpcServer; + } + + AdminSparkServer getAdminServer() { + return adminServer; + } + + VeniceGrpcServer getAdminGrpcServer() { + return adminGrpcServer; + } + + TopicCleanupService getTopicCleanupService() { + return topicCleanupService; + } + + Optional getStoreBackupVersionCleanupService() { + return storeBackupVersionCleanupService; + } + + Optional getDisabledPartitionEnablerService() { + return disabledPartitionEnablerService; + } + + Optional getUnusedValueSchemaCleanupService() { + return unusedValueSchemaCleanupService; + } + + Optional getStoreGraveyardCleanupService() { + return storeGraveyardCleanupService; + } } diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceControllerClusterConfig.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceControllerClusterConfig.java index 827c9024e37..ca516b90fcd 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceControllerClusterConfig.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceControllerClusterConfig.java @@ -26,6 +26,8 @@ import static com.linkedin.venice.ConfigKeys.CLUSTER_TO_D2; import static com.linkedin.venice.ConfigKeys.CLUSTER_TO_SERVER_D2; import static com.linkedin.venice.ConfigKeys.CONCURRENT_INIT_ROUTINES_ENABLED; +import static com.linkedin.venice.ConfigKeys.CONTROLLER_ADMIN_GRPC_PORT; +import static com.linkedin.venice.ConfigKeys.CONTROLLER_ADMIN_SECURE_GRPC_PORT; import static com.linkedin.venice.ConfigKeys.CONTROLLER_AUTO_MATERIALIZE_DAVINCI_PUSH_STATUS_SYSTEM_STORE; import static com.linkedin.venice.ConfigKeys.CONTROLLER_AUTO_MATERIALIZE_META_SYSTEM_STORE; import static com.linkedin.venice.ConfigKeys.CONTROLLER_BACKUP_VERSION_DEFAULT_RETENTION_MS; @@ -47,6 +49,8 @@ import static com.linkedin.venice.ConfigKeys.CONTROLLER_EARLY_DELETE_BACKUP_ENABLED; import static com.linkedin.venice.ConfigKeys.CONTROLLER_ENABLE_DISABLED_REPLICA_ENABLED; import static com.linkedin.venice.ConfigKeys.CONTROLLER_ENFORCE_SSL; +import static com.linkedin.venice.ConfigKeys.CONTROLLER_GRPC_SERVER_ENABLED; +import static com.linkedin.venice.ConfigKeys.CONTROLLER_GRPC_SERVER_THREAD_COUNT; import static com.linkedin.venice.ConfigKeys.CONTROLLER_HAAS_SUPER_CLUSTER_NAME; import static com.linkedin.venice.ConfigKeys.CONTROLLER_HELIX_CLOUD_ID; import static com.linkedin.venice.ConfigKeys.CONTROLLER_HELIX_CLOUD_INFO_PROCESSOR_NAME; @@ -177,6 +181,7 @@ import static com.linkedin.venice.utils.ByteUtils.BYTES_PER_MB; import static com.linkedin.venice.utils.ByteUtils.generateHumanReadableByteCountString; +import com.linkedin.venice.ConfigKeys; import com.linkedin.venice.PushJobCheckpoints; import com.linkedin.venice.SSLConfig; import com.linkedin.venice.authorization.DefaultIdentityParser; @@ -233,6 +238,8 @@ public class VeniceControllerClusterConfig { private final String adminHostname; private final int adminPort; private final int adminSecurePort; + private final int adminGrpcPort; + private final int adminSecureGrpcPort; private final int controllerClusterReplica; // Name of the Helix cluster for controllers private final String controllerClusterName; @@ -290,6 +297,8 @@ public class VeniceControllerClusterConfig { private final boolean backupVersionRetentionBasedCleanupEnabled; private final boolean backupVersionMetadataFetchBasedCleanupEnabled; + private final boolean grpcServerEnabled; + private final int grpcServerThreadCount; private final boolean enforceSSLOnly; private final long terminalStateTopicCheckerDelayMs; private final List disabledRoutes; @@ -534,6 +543,7 @@ public class VeniceControllerClusterConfig { private final long serviceDiscoveryRegistrationRetryMS; private Set pushJobUserErrorCheckpoints; + private boolean isHybridStorePartitionCountUpdateEnabled; public VeniceControllerClusterConfig(VeniceProperties props) { this.props = props; @@ -657,6 +667,12 @@ public VeniceControllerClusterConfig(VeniceProperties props) { this.adminPort = props.getInt(ADMIN_PORT); this.adminHostname = props.getString(ADMIN_HOSTNAME, Utils::getHostName); this.adminSecurePort = props.getInt(ADMIN_SECURE_PORT); + this.adminGrpcPort = props.getInt(CONTROLLER_ADMIN_GRPC_PORT, -1); + this.adminSecureGrpcPort = props.getInt(CONTROLLER_ADMIN_SECURE_GRPC_PORT, -1); + this.grpcServerEnabled = props.getBoolean(CONTROLLER_GRPC_SERVER_ENABLED, false); + this.grpcServerThreadCount = + props.getInt(CONTROLLER_GRPC_SERVER_THREAD_COUNT, Runtime.getRuntime().availableProcessors()); + /** * Override the config to false if the "Read" method check is not working as expected. */ @@ -854,8 +870,8 @@ public VeniceControllerClusterConfig(VeniceProperties props) { props.getBoolean(CONTROLLER_BACKUP_VERSION_RETENTION_BASED_CLEANUP_ENABLED, false); this.backupVersionMetadataFetchBasedCleanupEnabled = props.getBoolean(CONTROLLER_BACKUP_VERSION_METADATA_FETCH_BASED_CLEANUP_ENABLED, false); - this.enforceSSLOnly = props.getBoolean(CONTROLLER_ENFORCE_SSL, false); // By default, allow both secure and insecure - // routes + // By default, allow both secure and insecure routes + this.enforceSSLOnly = props.getBoolean(CONTROLLER_ENFORCE_SSL, false); this.terminalStateTopicCheckerDelayMs = props.getLong(TERMINAL_STATE_TOPIC_CHECK_DELAY_MS, TimeUnit.MINUTES.toMillis(10)); this.disableParentTopicTruncationUponCompletion = @@ -978,6 +994,8 @@ public VeniceControllerClusterConfig(VeniceProperties props) { this.serviceDiscoveryRegistrationRetryMS = props.getLong(SERVICE_DISCOVERY_REGISTRATION_RETRY_MS, 30L * Time.MS_PER_SECOND); this.pushJobUserErrorCheckpoints = parsePushJobUserErrorCheckpoints(props); + this.isHybridStorePartitionCountUpdateEnabled = + props.getBoolean(ConfigKeys.CONTROLLER_ENABLE_HYBRID_STORE_PARTITION_COUNT_UPDATE, false); } public VeniceProperties getProps() { @@ -1226,6 +1244,14 @@ public int getAdminSecurePort() { return adminSecurePort; } + public int getAdminGrpcPort() { + return adminGrpcPort; + } + + public int getAdminSecureGrpcPort() { + return adminSecureGrpcPort; + } + public boolean adminCheckReadMethodForKafka() { return adminCheckReadMethodForKafka; } @@ -1472,6 +1498,14 @@ public boolean isControllerEnforceSSLOnly() { return enforceSSLOnly; } + public boolean isGrpcServerEnabled() { + return grpcServerEnabled; + } + + public int getGrpcServerThreadCount() { + return grpcServerThreadCount; + } + public long getTerminalStateTopicCheckerDelayMs() { return terminalStateTopicCheckerDelayMs; } @@ -1767,6 +1801,10 @@ public int getDanglingTopicOccurrenceThresholdForCleanup() { return danglingTopicOccurrenceThresholdForCleanup; } + public boolean isHybridStorePartitionCountUpdateEnabled() { + return isHybridStorePartitionCountUpdateEnabled; + } + /** * A function that would put a k/v pair into a map with some processing works. */ diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceControllerMultiClusterConfig.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceControllerMultiClusterConfig.java index 2f9ea6d8a5b..0de71246bae 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceControllerMultiClusterConfig.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceControllerMultiClusterConfig.java @@ -55,6 +55,14 @@ public int getAdminSecurePort() { return getCommonConfig().getAdminSecurePort(); } + public int getAdminGrpcPort() { + return getCommonConfig().getAdminGrpcPort(); + } + + public int getAdminSecureGrpcPort() { + return getCommonConfig().getAdminSecureGrpcPort(); + } + public boolean adminCheckReadMethodForKafka() { return getCommonConfig().adminCheckReadMethodForKafka(); } @@ -219,6 +227,14 @@ public boolean isControllerEnforceSSLOnly() { return getCommonConfig().isControllerEnforceSSLOnly(); } + public boolean isGrpcServerEnabled() { + return getCommonConfig().isGrpcServerEnabled(); + } + + public int getGrpcServerThreadCount() { + return getCommonConfig().getGrpcServerThreadCount(); + } + public long getTerminalStateTopicCheckerDelayMs() { return getCommonConfig().getTerminalStateTopicCheckerDelayMs(); } diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceHelixAdmin.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceHelixAdmin.java index 8fc07ade4e2..1a8fb7703c9 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceHelixAdmin.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceHelixAdmin.java @@ -14,6 +14,7 @@ import static com.linkedin.venice.meta.Store.NON_EXISTING_VERSION; import static com.linkedin.venice.meta.Version.PushType; import static com.linkedin.venice.meta.VersionStatus.ERROR; +import static com.linkedin.venice.meta.VersionStatus.KILLED; import static com.linkedin.venice.meta.VersionStatus.NOT_CREATED; import static com.linkedin.venice.meta.VersionStatus.ONLINE; import static com.linkedin.venice.meta.VersionStatus.PUSHED; @@ -135,6 +136,7 @@ import com.linkedin.venice.meta.StoreDataAudit; import com.linkedin.venice.meta.StoreGraveyard; import com.linkedin.venice.meta.StoreInfo; +import com.linkedin.venice.meta.StoreName; import com.linkedin.venice.meta.SystemStoreAttributes; import com.linkedin.venice.meta.VeniceUserStoreType; import com.linkedin.venice.meta.Version; @@ -646,6 +648,7 @@ public VeniceHelixAdmin( // Participant stores are not read or written in parent colo. Parent controller skips participant store // initialization. if (!isParent() && multiClusterConfigs.isParticipantMessageStoreEnabled()) { + LOGGER.info("Adding PerClusterInternalRTStoreInitializationRoutine for ParticipantMessageStore"); initRoutines.add( new PerClusterInternalRTStoreInitializationRoutine( PARTICIPANT_MESSAGE_SYSTEM_STORE_VALUE, @@ -1864,7 +1867,7 @@ protected void checkPreConditionForCreateStore( String valueSchema, boolean allowSystemStore, boolean skipLingeringResourceCheck) { - if (!Store.isValidStoreName(storeName)) { + if (!StoreName.isValidStoreName(storeName)) { throw new VeniceException("Invalid store name " + storeName + ". Only letters, numbers, underscore or dash"); } AvroSchemaUtils.validateAvroSchemaStr(keySchema); @@ -2781,46 +2784,8 @@ private Pair addVersion( version.setPushStreamSourceAddress(sourceKafkaBootstrapServers); version.setNativeReplicationSourceFabric(sourceFabric); } - if (isParent() && ((store.isHybrid() - && store.getHybridStoreConfig().getDataReplicationPolicy() == DataReplicationPolicy.AGGREGATE) - || store.isIncrementalPushEnabled())) { - // Create rt topic in parent colo if the store is aggregate mode hybrid store - PubSubTopic realTimeTopic = pubSubTopicRepository.getTopic(Utils.getRealTimeTopicName(store)); - if (!getTopicManager().containsTopic(realTimeTopic)) { - getTopicManager().createTopic( - realTimeTopic, - numberOfPartitions, - clusterConfig.getKafkaReplicationFactorRTTopics(), - StoreUtils.getExpectedRetentionTimeInMs(store, store.getHybridStoreConfig()), - false, - // Note: do not enable RT compaction! Might make jobs in Online/Offline model stuck - clusterConfig.getMinInSyncReplicasRealTimeTopics(), - false); - if (version.isSeparateRealTimeTopicEnabled()) { - getTopicManager().createTopic( - pubSubTopicRepository.getTopic(Version.composeSeparateRealTimeTopic(storeName)), - numberOfPartitions, - clusterConfig.getKafkaReplicationFactorRTTopics(), - StoreUtils.getExpectedRetentionTimeInMs(store, store.getHybridStoreConfig()), - false, - // Note: do not enable RT compaction! Might make jobs in Online/Offline model stuck - clusterConfig.getMinInSyncReplicasRealTimeTopics(), - false); - } - } else { - // If real-time topic already exists, check whether its retention time is correct. - PubSubTopicConfiguration pubSubTopicConfiguration = - getTopicManager().getCachedTopicConfig(realTimeTopic); - long topicRetentionTimeInMs = TopicManager.getTopicRetention(pubSubTopicConfiguration); - long expectedRetentionTimeMs = - StoreUtils.getExpectedRetentionTimeInMs(store, store.getHybridStoreConfig()); - if (topicRetentionTimeInMs != expectedRetentionTimeMs) { - getTopicManager() - .updateTopicRetention(realTimeTopic, expectedRetentionTimeMs, pubSubTopicConfiguration); - } - } - } } + /** * Version-level rewind time override. */ @@ -2852,6 +2817,9 @@ private Pair addVersion( constructViewResources(veniceViewProperties, store, version.getNumber()); repository.updateStore(store); + if (isRealTimeTopicRequired(store, version)) { + createOrUpdateRealTimeTopics(clusterName, store, version); + } LOGGER.info("Add version: {} for store: {}", version.getNumber(), storeName); /** @@ -3010,6 +2978,172 @@ private Pair addVersion( } } + /** + * Determines whether real-time topics should be created for the given store and version. + * + *

Real-time topics are created based on the following conditions: + *

    + *
  • The store and version must both be hybrid ({@code store.isHybrid()} and {@code version.isHybrid()}).
    • + *
  • If the controller is a child, real-time topics are always created.
    • + *
  • If the controller is a parent, real-time topics are created only if: + *
      + *
    • Active-active replication is disabled for the store, and
      • + *
    • Either the store's data replication policy is {@code DataReplicationPolicy.AGGREGATE}, or
      • + *
    • Incremental push is enabled for the store.
      • + *
    + *
    • + *
+ * + * @param store the store being evaluated + * @param version the version being evaluated + * @return {@code true} if real-time topics should be created; {@code false} otherwise + */ + boolean isRealTimeTopicRequired(Store store, Version version) { + if (!store.isHybrid() || !version.isHybrid()) { + return false; + } + + // Child controllers always create real-time topics for hybrid stores in their region + if (!isParent()) { + return true; + } + + // Parent controllers create real-time topics in the parent region only under certain conditions + return !store.isActiveActiveReplicationEnabled() + && (store.getHybridStoreConfig().getDataReplicationPolicy() == DataReplicationPolicy.AGGREGATE + || store.isIncrementalPushEnabled()); + } + + /** + * Creates or updates real-time topics for the specified store (using reference hybrid version) in the given cluster. + * + *

This method ensures that real-time topics (primary and separate, if applicable) are configured + * correctly for a hybrid store. It creates the topics if they do not exist and updates their retention + * time if necessary. For stores with separate real-time topics enabled, the method handles the creation + * or update of those topics as well. + * + * @param clusterName the name of the cluster where the topics are managed + * @param store the {@link Store} associated with the topics + * @param version the {@link Version} containing the configuration for the topics, including partition count + * and hybrid store settings + */ + void createOrUpdateRealTimeTopics(String clusterName, Store store, Version version) { + LOGGER.info( + "Setting up real-time topics for store: {} with reference hybrid version: {} in cluster: {}", + store.getName(), + version.getNumber(), + clusterName); + String storeName = store.getName(); + // Create real-time topic if it doesn't exist; otherwise, update the retention time if necessary + PubSubTopic realTimeTopic = getPubSubTopicRepository().getTopic(Utils.getRealTimeTopicName(version)); + createOrUpdateRealTimeTopic(clusterName, store, version, realTimeTopic); + + // Create separate real-time topic if it doesn't exist; otherwise, update the retention time if necessary + if (version.isSeparateRealTimeTopicEnabled()) { + // TODO: Add support for repartitioning separate real-time topics, primarily needed for incremental push jobs. + createOrUpdateRealTimeTopic( + clusterName, + store, + version, + getPubSubTopicRepository().getTopic(Version.composeSeparateRealTimeTopic(storeName))); + } + } + + /** + * Creates or updates a real-time topic for a given store in the specified cluster. + * + *

This method ensures that the real-time topic matches the expected configuration based on + * the store's hybrid settings and the associated version. If the topic already exists: + *

    + *
  • It validates the partition count against the expected partition count for the version.
    • + *
  • It updates the retention time if necessary.
    • + *
+ * If the topic does not exist, it creates the topic with the required configuration. + * + * @param clusterName the name of the cluster where the topic resides + * @param store the {@link Store} store to which the topic belongs + * @param version the reference hybrid {@link Version} containing + * @param realTimeTopic the {@link PubSubTopic} representing the real-time topic + * @throws VeniceException if the partition count of an existing topic does not match the expected value + */ + void createOrUpdateRealTimeTopic(String clusterName, Store store, Version version, PubSubTopic realTimeTopic) { + int expectedNumOfPartitions = version.getPartitionCount(); + TopicManager topicManager = getTopicManager(); + if (topicManager.containsTopic(realTimeTopic)) { + validateAndUpdateTopic(realTimeTopic, store, version, expectedNumOfPartitions, topicManager); + } else { + VeniceControllerClusterConfig clusterConfig = getControllerConfig(clusterName); + topicManager.createTopic( + realTimeTopic, + expectedNumOfPartitions, + clusterConfig.getKafkaReplicationFactorRTTopics(), + StoreUtils.getExpectedRetentionTimeInMs(store, store.getHybridStoreConfig()), + false, + // Note: do not enable RT compaction! Might make jobs in Online/Offline model stuck + clusterConfig.getMinInSyncReplicasRealTimeTopics(), + false); + } + LOGGER.info( + "Completed setup for real-time topic: {} for store: {} with reference hybrid version: {} and partition count: {}", + realTimeTopic.getName(), + store.getName(), + version.getNumber(), + expectedNumOfPartitions); + } + + /** + * Validates the real-time topic's configuration and updates its retention time if necessary. + * + *

This method checks if the partition count of the real-time topic matches the expected partition count + * for the specified version. If the counts do not match, an exception is thrown. Additionally, it validates + * the topic's retention time against the expected retention time and updates it if required. + * + * @param realTimeTopic the {@link PubSubTopic} representing the real-time topic to validate + * @param store the {@link Store} store to which the topic belongs + * @param version the reference hybrid {@link Version} + * @param expectedNumOfPartitions the expected number of partitions for the real-time topic + * @param topicManager the {@link TopicManager} used for topic management operations + * @throws VeniceException if the partition count of the topic does not match the expected partition count + */ + void validateAndUpdateTopic( + PubSubTopic realTimeTopic, + Store store, + Version version, + int expectedNumOfPartitions, + TopicManager topicManager) { + int actualNumOfPartitions = topicManager.getPartitionCount(realTimeTopic); + // Validate partition count + if (actualNumOfPartitions != expectedNumOfPartitions) { + LOGGER.error( + "Real-time topic: {} for store: {} has different partition count: {} from version partition count: {} version: {} store: {}", + realTimeTopic.getName(), + store.getName(), + actualNumOfPartitions, + expectedNumOfPartitions, + version, + store); + String errorMessage = String.format( + "Real-time topic: %s for store: %s has different partition count: %d from version partition count: %d", + realTimeTopic.getName(), + store.getName(), + actualNumOfPartitions, + expectedNumOfPartitions); + throw new VeniceException(errorMessage); + } + + // Validate and update retention time if necessary + HybridStoreConfig hybridStoreConfig = store.getHybridStoreConfig(); + long expectedRetentionTimeMs = StoreUtils.getExpectedRetentionTimeInMs(store, hybridStoreConfig); + boolean isUpdated = topicManager.updateTopicRetentionWithRetries(realTimeTopic, expectedRetentionTimeMs); + LOGGER.info( + "{} retention time for real-time topic: {} for store: {} (hybrid version: {}, partition count: {})", + isUpdated ? "Updated" : "Validated", + realTimeTopic.getName(), + store.getName(), + version.getNumber(), + expectedNumOfPartitions); + } + /** * During store migration, skip a version if: * This is the child controller of the destination cluster @@ -3078,7 +3212,7 @@ public Version incrementVersionIdempotent( VeniceControllerClusterConfig clusterConfig = getHelixVeniceClusterResources(clusterName).getConfig(); int replicationMetadataVersionId = clusterConfig.getReplicationMetadataVersion(); return pushType.isIncremental() - ? getIncrementalPushVersion(clusterName, storeName) + ? getIncrementalPushVersion(clusterName, storeName, pushJobId) : addVersion( clusterName, storeName, @@ -3165,80 +3299,67 @@ private Optional getVersionWithPushId(String clusterName, String storeN } /** - * Get the real time topic name for a given store. If the topic is not created in Kafka, it creates the - * real time topic and returns the topic name. - * @param clusterName name of the Venice cluster. - * @param store store. - * @return name of the store's real time topic name. + * Ensures that a real-time topic exists for the given user system store. + * If the topic does not already exist in PubSub, it creates the real-time topic + * and returns the topic name. This method is specific to user system stores, + * where real-time topics are eagerly created by the controller to prevent + * blocking of threads that produce data to these stores. + * + * @param clusterName the name of the Venice cluster. + * @param storeName the name of the store. + * @return the name of the store's real-time topic. + * @throws VeniceNoStoreException if the store does not exist in the specified cluster. + * @throws VeniceException if the store is not a user system store or if the partition count is invalid. */ - @Override - public String getRealTimeTopic(String clusterName, Store store) { - checkControllerLeadershipFor(clusterName); - PubSubTopic realTimeTopic = pubSubTopicRepository.getTopic(Utils.getRealTimeTopicName(store)); - ensureRealTimeTopicIsReady(clusterName, realTimeTopic); - return realTimeTopic.getName(); - } - - @Override - public String getSeparateRealTimeTopic(String clusterName, String storeName) { + void ensureRealTimeTopicExistsForUserSystemStores(String clusterName, String storeName) { checkControllerLeadershipFor(clusterName); - PubSubTopic incrementalPushRealTimeTopic = - pubSubTopicRepository.getTopic(Version.composeSeparateRealTimeTopic(storeName)); - ensureRealTimeTopicIsReady(clusterName, incrementalPushRealTimeTopic); - return incrementalPushRealTimeTopic.getName(); - } - - private void ensureRealTimeTopicIsReady(String clusterName, PubSubTopic realTimeTopic) { + Store store = getStore(clusterName, storeName); + if (store == null) { + throw new VeniceNoStoreException(storeName, clusterName); + } + VeniceSystemStoreType systemStoreType = VeniceSystemStoreType.getSystemStoreType(storeName); + if (VeniceSystemStoreType.META_STORE != systemStoreType + && VeniceSystemStoreType.DAVINCI_PUSH_STATUS_STORE != systemStoreType) { + LOGGER.error("Failed to create real time topic for store: {} because it is not a user system store.", storeName); + throw new VeniceException( + "Failed to create real time topic for store: " + storeName + " because it is not a user system store."); + } + PubSubTopic realTimeTopic = getPubSubTopicRepository().getTopic(Utils.getRealTimeTopicName(store)); TopicManager topicManager = getTopicManager(); - String storeName = realTimeTopic.getStoreName(); - if (!topicManager.containsTopic(realTimeTopic)) { - HelixVeniceClusterResources resources = getHelixVeniceClusterResources(clusterName); - try (AutoCloseableLock ignore = resources.getClusterLockManager().createStoreWriteLock(storeName)) { - // The topic might be created by another thread already. Check before creating. - if (topicManager.containsTopic(realTimeTopic)) { - return; - } - ReadWriteStoreRepository repository = resources.getStoreMetadataRepository(); - Store store = repository.getStore(storeName); - if (store == null) { - throwStoreDoesNotExist(clusterName, storeName); - } - if (!store.isHybrid() && !store.isWriteComputationEnabled() && !store.isSystemStore()) { - logAndThrow("Store " + storeName + " is not hybrid, refusing to return a realtime topic"); - } - Version version = store.getVersion(store.getLargestUsedVersionNumber()); - int partitionCount = version != null ? version.getPartitionCount() : 0; - // during transition to version based partition count, some old stores may have partition count on the store - // config only. - if (partitionCount == 0) { - // Now store-level partition count is set when a store is converted to hybrid - partitionCount = store.getPartitionCount(); - if (partitionCount == 0) { - if (version == null) { - throw new VeniceException("Store: " + storeName + " is not initialized with a version yet"); - } else { - throw new VeniceException("Store: " + storeName + " has partition count set to 0"); - } - } - } + if (topicManager.containsTopic(realTimeTopic)) { + return; + } - VeniceControllerClusterConfig clusterConfig = getHelixVeniceClusterResources(clusterName).getConfig(); - getTopicManager().createTopic( - realTimeTopic, - partitionCount, - clusterConfig.getKafkaReplicationFactorRTTopics(), - store.getRetentionTime(), - false, - // Note: do not enable RT compaction! Might make jobs in Online/Offline model stuck - clusterConfig.getMinInSyncReplicasRealTimeTopics(), - false); - // TODO: if there is an online version from a batch push before this store was hybrid then we won't start - // replicating to it. A new version must be created. - LOGGER.warn( - "Creating real time topic per topic request for store: {}. " - + "Buffer replay won't start for any existing versions", + HelixVeniceClusterResources resources = getHelixVeniceClusterResources(clusterName); + try (AutoCloseableLock ignore = resources.getClusterLockManager().createStoreWriteLock(storeName)) { + // check again that the real-time topic does not exist + if (topicManager.containsTopic(realTimeTopic)) { + return; + } + Version version = store.getVersion(store.getLargestUsedVersionNumber()); + int partitionCount = version != null ? version.getPartitionCount() : store.getPartitionCount(); + if (partitionCount == 0) { + LOGGER.error( + "Failed to create real time topic for user system store: {} because both store and version have partition count set to 0.", storeName); + throw new VeniceException( + "Failed to create real time topic for user system store: " + storeName + + " because both store and version have partition count set to 0."); } + VeniceControllerClusterConfig clusterConfig = getControllerConfig(clusterName); + LOGGER.info( + "Creating real time topic for user system store: {} with partition count: {}", + storeName, + partitionCount); + getTopicManager().createTopic( + realTimeTopic, + partitionCount, + clusterConfig.getKafkaReplicationFactorRTTopics(), + store.getRetentionTime(), + false, + // Note: do not enable RT compaction! Might make jobs in Online/Offline model stuck + clusterConfig.getMinInSyncReplicasRealTimeTopics(), + false); } } @@ -3261,47 +3382,148 @@ public Optional getReplicationMetadataSchema( } } - /** - * @see Admin#getIncrementalPushVersion(String, String) - */ @Override - public Version getIncrementalPushVersion(String clusterName, String storeName) { + public Version getReferenceVersionForStreamingWrites(String clusterName, String storeName, String pushJobId) { checkControllerLeadershipFor(clusterName); HelixVeniceClusterResources resources = getHelixVeniceClusterResources(clusterName); try (AutoCloseableLock ignore = resources.getClusterLockManager().createStoreReadLock(storeName)) { + validateStoreSetupForRTWrites(clusterName, storeName, pushJobId, PushType.STREAM); Store store = resources.getStoreMetadataRepository().getStore(storeName); - if (store == null) { - throwStoreDoesNotExist(clusterName, storeName); + Version hybridVersion = getReferenceHybridVersionForRealTimeWrites(clusterName, store, pushJobId); + if (!isParent()) { + PubSubTopic rtTopic = getPubSubTopicRepository().getTopic(Utils.getRealTimeTopicName(hybridVersion)); + int partitionCount = hybridVersion.getPartitionCount(); + validateTopicPresenceAndState(clusterName, storeName, pushJobId, PushType.STREAM, rtTopic, partitionCount); } + return hybridVersion; + } + } - if (!store.isIncrementalPushEnabled()) { - throw new VeniceException("Incremental push is not enabled for store: " + storeName); + @Override + public Version getIncrementalPushVersion(String clusterName, String storeName, String pushJobId) { + checkControllerLeadershipFor(clusterName); + HelixVeniceClusterResources resources = getHelixVeniceClusterResources(clusterName); + try (AutoCloseableLock ignore = resources.getClusterLockManager().createStoreReadLock(storeName)) { + validateStoreSetupForRTWrites(clusterName, storeName, pushJobId, PushType.INCREMENTAL); + Store store = resources.getStoreMetadataRepository().getStore(storeName); + Version hybridVersion = getReferenceHybridVersionForRealTimeWrites(clusterName, store, pushJobId); + // If real-time topic is required, validate that it exists and is in a good state + if (isRealTimeTopicRequired(store, hybridVersion)) { + validateTopicForIncrementalPush(clusterName, store, hybridVersion, pushJobId); } + return hybridVersion; + } + } - List versions = store.getVersions(); - if (versions.isEmpty()) { - throw new VeniceException("Store: " + storeName + " is not initialized with a version yet"); - } + void validateStoreSetupForRTWrites(String clusterName, String storeName, String pushJobId, PushType pushType) { + Store store = getHelixVeniceClusterResources(clusterName).getStoreMetadataRepository().getStore(storeName); + if (store == null) { + throwStoreDoesNotExist(clusterName, storeName); + } + if (!store.isHybrid()) { + LOGGER.error( + "{} push writes with pushJobId: {} on store: {} in cluster: {} are not allowed because it is not a hybrid store", + pushType, + pushJobId, + storeName, + clusterName); + throw new VeniceException( + "Store: " + storeName + " is not a hybrid store and cannot be used for " + pushType + " writes"); + } + if (pushType == PushType.INCREMENTAL && !store.isIncrementalPushEnabled()) { + LOGGER.error( + "Incremental push with pushJobId: {} on store: {} in cluster: {} is not allowed because incremental push is not enabled", + pushJobId, + storeName, + clusterName); + throw new VeniceException("Store: " + storeName + " is not an incremental push store"); + } + } - /** - * Don't use {@link Store#getCurrentVersion()} here since it is always 0 in parent controller - */ - Version version = versions.get(versions.size() - 1); - if (version.getStatus() == ERROR) { - throw new VeniceException( - "cannot have incremental push because current version is in error status. " + "Version: " - + version.getNumber() + " Store:" + storeName); - } + void validateTopicForIncrementalPush( + String clusterName, + Store store, + Version referenceHybridVersion, + String pushJobId) { + PubSubTopicRepository topicRepository = getPubSubTopicRepository(); + if (referenceHybridVersion.isSeparateRealTimeTopicEnabled()) { + PubSubTopic separateRtTopic = topicRepository.getTopic(Version.composeSeparateRealTimeTopic(store.getName())); + validateTopicPresenceAndState( + clusterName, + store.getName(), + pushJobId, + PushType.INCREMENTAL, + separateRtTopic, + referenceHybridVersion.getPartitionCount()); + // We can consider short-circuiting here if the separate real-time topic is enabled and + // the topic is in a good state + } - PubSubTopic rtTopic = pubSubTopicRepository.getTopic(Utils.getRealTimeTopicName(store)); - if (!getTopicManager().containsTopicAndAllPartitionsAreOnline(rtTopic) || isTopicTruncated(rtTopic.getName())) { - resources.getVeniceAdminStats().recordUnexpectedTopicAbsenceCount(); - throw new VeniceException( - "Incremental push cannot be started for store: " + storeName + " in cluster: " + clusterName - + " because the topic: " + rtTopic + " is either absent or being truncated"); + PubSubTopic rtTopic = topicRepository.getTopic(Utils.getRealTimeTopicName(referenceHybridVersion)); + validateTopicPresenceAndState( + clusterName, + store.getName(), + pushJobId, + PushType.INCREMENTAL, + rtTopic, + referenceHybridVersion.getPartitionCount()); + } + + void validateTopicPresenceAndState( + String clusterName, + String storeName, + String pushJobId, + PushType pushType, + PubSubTopic topic, + int partitionCount) { + if (getTopicManager().containsTopicAndAllPartitionsAreOnline(topic, partitionCount) + && !isTopicTruncated(topic.getName())) { + return; + } + LOGGER.error( + "{} push writes from pushJobId: {} cannot be accepted on store: {} in cluster: {} because the topic: {} is either absent or being truncated", + pushType, + pushJobId, + storeName, + clusterName, + topic); + getHelixVeniceClusterResources(clusterName).getVeniceAdminStats().recordUnexpectedTopicAbsenceCount(); + throw new VeniceException( + pushType + " push writes cannot be accepted on store: " + storeName + " in cluster: " + clusterName + + " because the topic: " + topic + " is either absent or being truncated"); + } + + Version getReferenceHybridVersionForRealTimeWrites(String clusterName, Store store, String pushJobId) { + List versions = new ArrayList<>(store.getVersions()); + if (versions.isEmpty()) { + LOGGER.error( + "Store: {} in cluster: {} is not initialized with a version yet. Rejecting request for writes with pushJobId: {}", + store.getName(), + clusterName, + pushJobId); + throw new VeniceException("Store: " + store.getName() + " is not initialized with a version yet."); + } + + versions.sort(Comparator.comparingInt(Version::getNumber).reversed()); + for (Version version: versions) { + if (version.getHybridStoreConfig() != null && version.getStatus() != ERROR && version.getStatus() != KILLED) { + LOGGER.info( + "Found hybrid version: {} for store: {} in cluster: {}. Will use it as a reference for pushJobId: {}", + version.getNumber(), + version, + clusterName, + pushJobId); + return version; } - return version; } + + String logMessage = String.format( + "No valid hybrid store version (non-errored) found in store: %s in cluster: %s for pushJobId: %s.", + store.getName(), + clusterName, + pushJobId); + LOGGER.error(logMessage); + throw new VeniceException(logMessage); } /** @@ -3522,7 +3744,7 @@ private void deleteOneStoreVersion(String clusterName, String storeName, int ver } cleanUpViewResources(new Properties(), store, deletedVersion.get().getNumber()); } - if (store.isDaVinciPushStatusStoreEnabled()) { + if (store.isDaVinciPushStatusStoreEnabled() && !isParent()) { ExecutorService executor = Executors.newSingleThreadExecutor(); Future future = executor.submit( () -> getPushStatusStoreWriter().deletePushStatus( @@ -3890,6 +4112,7 @@ private boolean truncateKafkaTopic( String kafkaTopicName, long deprecatedJobTopicRetentionMs) { try { + if (topicManager .updateTopicRetention(pubSubTopicRepository.getTopic(kafkaTopicName), deprecatedJobTopicRetentionMs)) { return true; @@ -4275,8 +4498,13 @@ public void setStorePartitionCount(String clusterName, String storeName, int par preCheckStorePartitionCountUpdate(clusterName, store, partitionCount); // Do not update the partitionCount on the store.version as version config is immutable. The // version.getPartitionCount() - // is read only in getRealTimeTopic and createInternalStore creation, so modifying currentVersion should not have + // is read only in ensureRealTimeTopicExistsForUserSystemStores and createInternalStore creation, so modifying + // currentVersion should not have // any effect. + if (store.isHybrid() + && multiClusterConfigs.getControllerConfig(clusterName).isHybridStorePartitionCountUpdateEnabled()) { + generateAndUpdateRealTimeTopicName(store); + } if (partitionCount != 0) { store.setPartitionCount(partitionCount); } else { @@ -4290,7 +4518,16 @@ public void setStorePartitionCount(String clusterName, String storeName, int par void preCheckStorePartitionCountUpdate(String clusterName, Store store, int newPartitionCount) { String errorMessagePrefix = "Store update error for " + store.getName() + " in cluster: " + clusterName + ": "; VeniceControllerClusterConfig clusterConfig = getHelixVeniceClusterResources(clusterName).getConfig(); + int maxPartitionNum = clusterConfig.getMaxNumberOfPartitions(); + if (store.isHybrid() && store.getPartitionCount() != newPartitionCount) { + if (multiClusterConfigs.getControllerConfig(clusterName).isHybridStorePartitionCountUpdateEnabled() + && newPartitionCount <= maxPartitionNum && newPartitionCount >= 0) { + LOGGER.info( + "Allow updating store " + store.getName() + " partition count to " + newPartitionCount + + " because `updateRealTimeTopic` is true."); + return; + } // Allow the update if partition count is not configured and the new partition count matches RT partition count if (store.getPartitionCount() == 0) { TopicManager topicManager; @@ -4313,7 +4550,6 @@ void preCheckStorePartitionCountUpdate(String clusterName, Store store, int newP throw new VeniceHttpException(HttpStatus.SC_BAD_REQUEST, errorMessage, ErrorType.INVALID_CONFIG); } - int maxPartitionNum = clusterConfig.getMaxNumberOfPartitions(); if (newPartitionCount > maxPartitionNum) { String errorMessage = errorMessagePrefix + "Partition count: " + newPartitionCount + " should be less than max: " + maxPartitionNum; @@ -4327,6 +4563,20 @@ void preCheckStorePartitionCountUpdate(String clusterName, Store store, int newP } } + private void generateAndUpdateRealTimeTopicName(Store store) { + // get oldRealTimeTopicName from the store config because that will be more (or equally) recent than any version + // config + String oldRealTimeTopicName = Utils.getRealTimeTopicNameFromStoreConfig(store); + String newRealTimeTopicName = Utils.createNewRealTimeTopicName(oldRealTimeTopicName); + PubSubTopic newRealTimeTopic = getPubSubTopicRepository().getTopic(newRealTimeTopicName); + + if (getTopicManager().containsTopic(newRealTimeTopic)) { + throw new VeniceException("Topic " + newRealTimeTopic + " should not exist."); + } + + store.getHybridStoreConfig().setRealTimeTopicName(newRealTimeTopicName); + } + void setStorePartitionerConfig(String clusterName, String storeName, PartitionerConfig partitionerConfig) { storeMetadataUpdate(clusterName, storeName, store -> { // Only amplification factor is allowed to be changed if the store is a hybrid store. @@ -4804,7 +5054,8 @@ private void internalUpdateStore(String clusterName, String storeName, UpdateSto Optional hybridTimeLagThreshold = params.getHybridTimeLagThreshold(); Optional hybridDataReplicationPolicy = params.getHybridDataReplicationPolicy(); Optional hybridBufferReplayPolicy = params.getHybridBufferReplayPolicy(); - Optional realTimeTopicName = params.getRealTimeTopicName(); + Optional realTimeTopicName = Optional.empty(); // real time topic name should only be changed during + // partition count update Optional accessControlled = params.getAccessControlled(); Optional compressionStrategy = params.getCompressionStrategy(); Optional clientDecompressionEnabled = params.getClientDecompressionEnabled(); @@ -6166,7 +6417,7 @@ public OfflinePushStatusInfo getOffLinePushStatus( // if status is not SOIP remove incremental push version from the supposedlyOngoingIncrementalPushVersions if (incrementalPushVersion.isPresent() && (status == ExecutionStatus.END_OF_INCREMENTAL_PUSH_RECEIVED || status == ExecutionStatus.NOT_CREATED) - && store.isDaVinciPushStatusStoreEnabled()) { + && store.isDaVinciPushStatusStoreEnabled() && !isParent()) { getPushStatusStoreWriter().removeFromSupposedlyOngoingIncrementalPushVersions( store.getName(), versionNumber, @@ -6778,7 +7029,9 @@ public Instance getLeaderController(String clusterName) { id, Utils.parseHostFromHelixNodeIdentifier(id), Utils.parsePortFromHelixNodeIdentifier(id), - multiClusterConfigs.getAdminSecurePort()); + multiClusterConfigs.getAdminSecurePort(), + multiClusterConfigs.getAdminGrpcPort(), + multiClusterConfigs.getAdminSecureGrpcPort()); } } if (attempt < maxAttempts) { @@ -7356,7 +7609,7 @@ public void close() { void checkControllerLeadershipFor(String clusterName) { if (!isLeaderControllerFor(clusterName)) { throw new VeniceException( - "This controller:" + controllerName + " is not the leader controller for " + clusterName); + "This controller:" + controllerName + " is not the leader controller for cluster: " + clusterName); } } @@ -8022,7 +8275,11 @@ private void setUpDaVinciPushStatusStore(String clusterName, String storeName) { throwStoreDoesNotExist(clusterName, storeName); } String daVinciPushStatusStoreName = VeniceSystemStoreType.DAVINCI_PUSH_STATUS_STORE.getSystemStoreName(storeName); - getRealTimeTopic(clusterName, daVinciPushStatusStoreName); + + if (!isParent()) { + // We do not materialize PS3 for parent region. Hence, skip RT topic creation. + ensureRealTimeTopicExistsForUserSystemStores(clusterName, daVinciPushStatusStoreName); + } if (!store.isDaVinciPushStatusStoreEnabled()) { storeMetadataUpdate(clusterName, storeName, (s) -> { s.setDaVinciPushStatusStoreEnabled(true); @@ -8031,6 +8288,11 @@ private void setUpDaVinciPushStatusStore(String clusterName, String storeName) { } } + /** + * Set up the meta store and produce snapshot to meta store RT. Should be called in the child controllers. + * @param clusterName The cluster name. + * @param regularStoreName The regular user store name. + */ void setUpMetaStoreAndMayProduceSnapshot(String clusterName, String regularStoreName) { checkControllerLeadershipFor(clusterName); ReadWriteStoreRepository repository = getHelixVeniceClusterResources(clusterName).getStoreMetadataRepository(); @@ -8039,9 +8301,13 @@ void setUpMetaStoreAndMayProduceSnapshot(String clusterName, String regularStore throwStoreDoesNotExist(clusterName, regularStoreName); } - // Make sure RT topic exists before producing. There's no write to parent region meta store RT, but we still create - // the RT topic to be consistent in case it was not auto-materialized - getRealTimeTopic(clusterName, VeniceSystemStoreType.META_STORE.getSystemStoreName(regularStoreName)); + // Make sure RT topic exists before producing. + if (!isParent()) { + // We do not materialize meta store for parent region. Hence, skip RT topic creation. + ensureRealTimeTopicExistsForUserSystemStores( + clusterName, + VeniceSystemStoreType.META_STORE.getSystemStoreName(regularStoreName)); + } // Update the store flag to enable meta system store. if (!store.isStoreMetaSystemStoreEnabled()) { @@ -8420,6 +8686,9 @@ public PushStatusStoreWriter getPushStatusStoreWriter() { @Override public void sendHeartbeatToSystemStore(String clusterName, String storeName, long heartbeatTimeStamp) { + if (isParent()) { + return; + } VeniceSystemStoreType systemStoreType = VeniceSystemStoreType.getSystemStoreType(storeName); String userStoreName = systemStoreType.extractRegularStoreName(storeName); long currentTimestamp = System.currentTimeMillis(); @@ -8463,6 +8732,10 @@ VeniceControllerMultiClusterConfig getMultiClusterConfigs() { return multiClusterConfigs; } + VeniceControllerClusterConfig getControllerConfig(String clusterName) { + return multiClusterConfigs.getControllerConfig(clusterName); + } + // Only for testing public void setPushJobDetailsStoreClient(AvroSpecificStoreClient client) { pushJobDetailsStoreClient = client; diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceParentHelixAdmin.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceParentHelixAdmin.java index 9bd42282188..4f333f28ba4 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceParentHelixAdmin.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/VeniceParentHelixAdmin.java @@ -1532,7 +1532,7 @@ public Version incrementVersionIdempotent( Version newVersion; if (pushType.isIncremental()) { - newVersion = getVeniceHelixAdmin().getIncrementalPushVersion(clusterName, storeName); + newVersion = getVeniceHelixAdmin().getIncrementalPushVersion(clusterName, storeName, pushJobId); } else { validateTargetedRegions(targetedRegions, clusterName); @@ -1709,19 +1709,6 @@ private AddVersion getAddVersionMessage( return addVersion; } - /** - * @see VeniceHelixAdmin#getRealTimeTopic(String, Store) - */ - @Override - public String getRealTimeTopic(String clusterName, Store store) { - return getVeniceHelixAdmin().getRealTimeTopic(clusterName, store); - } - - @Override - public String getSeparateRealTimeTopic(String clusterName, String storeName) { - return getVeniceHelixAdmin().getSeparateRealTimeTopic(clusterName, storeName); - } - /** * A couple of extra checks are needed in parent controller * 1. check batch job statuses across child controllers. (We cannot only check the version status @@ -1730,14 +1717,19 @@ public String getSeparateRealTimeTopic(String clusterName, String storeName) { * preserve incremental push topic in parent Kafka anymore */ @Override - public Version getIncrementalPushVersion(String clusterName, String storeName) { - Version incrementalPushVersion = getVeniceHelixAdmin().getIncrementalPushVersion(clusterName, storeName); + public Version getIncrementalPushVersion(String clusterName, String storeName, String pushJobId) { + Version incrementalPushVersion = getVeniceHelixAdmin().getIncrementalPushVersion(clusterName, storeName, pushJobId); String incrementalPushTopic = incrementalPushVersion.kafkaTopicName(); ExecutionStatus status = getOffLinePushStatus(clusterName, incrementalPushTopic).getExecutionStatus(); return getIncrementalPushVersion(incrementalPushVersion, status); } + @Override + public Version getReferenceVersionForStreamingWrites(String clusterName, String storeName, String pushJobId) { + return getVeniceHelixAdmin().getReferenceVersionForStreamingWrites(clusterName, storeName, pushJobId); + } + // This method is only for internal / test use case Version getIncrementalPushVersion(Version incrementalPushVersion, ExecutionStatus status) { String storeName = incrementalPushVersion.getStoreName(); diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/init/SystemStoreInitializationHelper.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/init/SystemStoreInitializationHelper.java index 60b2589ef02..9dca7d2af3a 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/init/SystemStoreInitializationHelper.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/init/SystemStoreInitializationHelper.java @@ -64,6 +64,7 @@ public static void setupSystemStore( UpdateStoreQueryParams updateStoreQueryParams, Admin admin, VeniceControllerMultiClusterConfig multiClusterConfigs) { + LOGGER.info("Setting up system store: {} in cluster: {}", systemStoreName, clusterName); Map protocolSchemaMap = Utils.getAllSchemasFromResources(protocolDefinition); Store store = admin.getStore(clusterName, systemStoreName); String keySchemaString = keySchema != null ? keySchema.toString() : DEFAULT_KEY_SCHEMA_STR; @@ -86,7 +87,7 @@ public static void setupSystemStore( throw new VeniceException("Unable to create or fetch store " + systemStoreName); } } else { - LOGGER.info("Internal store {} already exists in cluster {}", systemStoreName, clusterName); + LOGGER.info("Internal store: {} already exists in cluster: {}", systemStoreName, clusterName); if (keySchema != null) { /** * Only verify the key schema if it is explicitly specified by the caller, and we don't care @@ -203,6 +204,8 @@ public static void setupSystemStore( LOGGER.info("Created a version for internal store {} in cluster {}", systemStoreName, clusterName); } + + LOGGER.info("System store: {} in cluster: {} is set up", systemStoreName, clusterName); } // Visible for testing diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumerService.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumerService.java index 32d48435921..abceff1af2e 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumerService.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumerService.java @@ -109,7 +109,6 @@ private AdminConsumptionTask getAdminConsumptionTaskForCluster(String clusterNam config.getAdminConsumptionCycleTimeoutMs(), config.getAdminConsumptionMaxWorkerThreadPoolSize(), pubSubTopicRepository, - pubSubMessageDeserializer, config.getRegionName()); } diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTask.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTask.java index 4b7299f7f2b..c294fc1a4e2 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTask.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTask.java @@ -26,7 +26,6 @@ import com.linkedin.venice.pubsub.PubSubTopicRepository; import com.linkedin.venice.pubsub.api.PubSubConsumerAdapter; import com.linkedin.venice.pubsub.api.PubSubMessage; -import com.linkedin.venice.pubsub.api.PubSubMessageDeserializer; import com.linkedin.venice.pubsub.api.PubSubTopic; import com.linkedin.venice.pubsub.api.PubSubTopicPartition; import com.linkedin.venice.pubsub.manager.TopicManager; @@ -38,12 +37,14 @@ import java.io.Closeable; import java.io.IOException; import java.util.ArrayList; +import java.util.HashSet; import java.util.Iterator; import java.util.LinkedList; import java.util.List; import java.util.Map; import java.util.Optional; import java.util.Queue; +import java.util.Set; import java.util.concurrent.Callable; import java.util.concurrent.CancellationException; import java.util.concurrent.ConcurrentHashMap; @@ -154,13 +155,11 @@ public String toString() { private volatile long failingOffset = UNASSIGNED_VALUE; private boolean topicExists; /** - * A {@link Map} of stores to admin operations belonging to each store. The corresponding kafka offset of each admin - * operation is also attached as the first element of the {@link Pair}. + * A {@link Map} of stores to admin operations belonging to each store. The corresponding kafka offset and other + * metadata of each admin operation are included in the {@link AdminOperationWrapper}. */ private final Map> storeAdminOperationsMapWithOffset; - private final ConcurrentHashMap storeToScheduledTask; - /** * Map of store names that have encountered some sort of exception during consumption to {@link AdminErrorInfo} * that has the details about the exception and the offset of the problematic admin message. @@ -169,7 +168,7 @@ public String toString() { private final Queue> undelegatedRecords; private final ExecutionIdAccessor executionIdAccessor; - private final ExecutorService executorService; + private ExecutorService executorService; private TopicManager sourceKafkaClusterTopicManager; @@ -239,10 +238,6 @@ public ExecutorService getExecutorService() { */ private long lastUpdateTimeForConsumptionOffsetLag = 0; - private final PubSubTopicRepository pubSubTopicRepository; - - private final PubSubMessageDeserializer pubSubMessageDeserializer; - /** * The local region name of the controller. */ @@ -263,7 +258,6 @@ public AdminConsumptionTask( long processingCycleTimeoutInMs, int maxWorkerThreadPoolSize, PubSubTopicRepository pubSubTopicRepository, - PubSubMessageDeserializer pubSubMessageDeserializer, String regionName) { this.clusterName = clusterName; this.topic = AdminTopicUtils.getTopicNameFromClusterName(clusterName); @@ -286,7 +280,6 @@ public AdminConsumptionTask( this.storeAdminOperationsMapWithOffset = new ConcurrentHashMap<>(); this.problematicStores = new ConcurrentHashMap<>(); - this.storeToScheduledTask = new ConcurrentHashMap<>(); // since we use an unbounded queue the core pool size is really the max pool size this.executorService = new ThreadPoolExecutor( maxWorkerThreadPoolSize, @@ -297,8 +290,6 @@ public AdminConsumptionTask( new DaemonThreadFactory(String.format("Venice-Admin-Execution-Task-%s", clusterName))); this.undelegatedRecords = new LinkedList<>(); this.stats.setAdminConsumptionFailedOffset(failingOffset); - this.pubSubTopicRepository = pubSubTopicRepository; - this.pubSubMessageDeserializer = pubSubMessageDeserializer; this.pubSubTopic = pubSubTopicRepository.getTopic(topic); this.regionName = regionName; @@ -311,6 +302,11 @@ public AdminConsumptionTask( } } + // For testing purpose only + void setAdminExecutionTaskExecutorService(ExecutorService executorService) { + this.executorService = executorService; + } + @Override public synchronized void close() throws IOException { isRunning.getAndSet(false); @@ -376,28 +372,29 @@ public void run() { } while (!undelegatedRecords.isEmpty()) { + PubSubMessage record = undelegatedRecords.peek(); + if (record == null) { + break; + } try { - long executionId = delegateMessage(undelegatedRecords.peek()); + long executionId = delegateMessage(record); if (executionId == lastDelegatedExecutionId) { - updateLastOffset(undelegatedRecords.peek().getOffset()); + updateLastOffset(record.getOffset()); } undelegatedRecords.remove(); } catch (DataValidationException dve) { // Very unlikely but DataValidationException could be thrown here. LOGGER.error( "Admin consumption task is blocked due to DataValidationException with offset {}", - undelegatedRecords.peek().getOffset(), + record.getOffset(), dve); - failingOffset = undelegatedRecords.peek().getOffset(); + failingOffset = record.getOffset(); stats.recordFailedAdminConsumption(); stats.recordAdminTopicDIVErrorReportCount(); break; } catch (Exception e) { - LOGGER.error( - "Admin consumption task is blocked due to Exception with offset {}", - undelegatedRecords.peek().getOffset(), - e); - failingOffset = undelegatedRecords.peek().getOffset(); + LOGGER.error("Admin consumption task is blocked due to Exception with offset {}", record.getOffset(), e); + failingOffset = record.getOffset(); stats.recordFailedAdminConsumption(); break; } @@ -461,7 +458,6 @@ private void unSubscribe() { storeAdminOperationsMapWithOffset.clear(); problematicStores.clear(); undelegatedRecords.clear(); - storeToScheduledTask.clear(); failingOffset = UNASSIGNED_VALUE; offsetToSkip = UNASSIGNED_VALUE; offsetToSkipDIV = UNASSIGNED_VALUE; @@ -482,6 +478,8 @@ private void unSubscribe() { } /** + * Package private for testing purpose + * * Delegate work from the {@code storeAdminOperationsMapWithOffset} to the worker threads. Wait for the worker threads * to complete or when timeout {@code processingCycleTimeoutInMs} is reached. Collect the result of each thread. * The result can either be success: all given {@link AdminOperation}s were processed successfully or made progress @@ -492,39 +490,51 @@ private void unSubscribe() { private void executeMessagesAndCollectResults() throws InterruptedException { lastSucceededExecutionIdMap = new ConcurrentHashMap<>(executionIdAccessor.getLastSucceededExecutionIdMap(clusterName)); + /** This set is used to track which store has a task scheduled, so that we schedule at most one per store. */ + Set storesWithScheduledTask = new HashSet<>(); + /** List of tasks to be executed by the worker threads. */ List> tasks = new ArrayList<>(); + /** + * Note that tasks and stores are parallel lists (the elements at each index correspond to one another), and both + * lists are also parallel to the results list, declared later in this function. + */ List stores = new ArrayList<>(); // Create a task for each store that has admin messages pending to be processed. boolean skipOffsetCommandHasBeenProcessed = false; for (Map.Entry> entry: storeAdminOperationsMapWithOffset.entrySet()) { - if (!entry.getValue().isEmpty()) { - long adminMessageOffset = entry.getValue().peek().getOffset(); - if (checkOffsetToSkip(adminMessageOffset, false)) { - entry.getValue().remove(); + String storeName = entry.getKey(); + Queue storeQueue = entry.getValue(); + if (!storeQueue.isEmpty()) { + AdminOperationWrapper nextOp = storeQueue.peek(); + if (nextOp == null) { + continue; + } + long adminMessageOffset = nextOp.getOffset(); + if (checkOffsetToSkip(nextOp.getOffset(), false)) { + storeQueue.remove(); skipOffsetCommandHasBeenProcessed = true; } AdminExecutionTask newTask = new AdminExecutionTask( LOGGER, clusterName, - entry.getKey(), + storeName, lastSucceededExecutionIdMap, lastPersistedExecutionId, - entry.getValue(), + storeQueue, admin, executionIdAccessor, isParentController, stats, - regionName, - storeToScheduledTask); + regionName); // Check if there is previously created scheduled task still occupying one thread from the pool. - if (storeToScheduledTask.putIfAbsent(entry.getKey(), newTask) == null) { + if (storesWithScheduledTask.add(storeName)) { // Log the store name and the offset of the task being added into the task list LOGGER.info( "Adding admin message from store {} with offset {} to the task list", - entry.getKey(), + storeName, adminMessageOffset); tasks.add(newTask); - stores.add(entry.getKey()); + stores.add(storeName); } } } @@ -550,21 +560,24 @@ private void executeMessagesAndCollectResults() throws InterruptedException { try { result.get(); problematicStores.remove(storeName); - if (internalQueuesEmptied && storeAdminOperationsMapWithOffset.containsKey(storeName) - && !storeAdminOperationsMapWithOffset.get(storeName).isEmpty()) { - internalQueuesEmptied = false; + if (internalQueuesEmptied) { + Queue storeQueue = storeAdminOperationsMapWithOffset.get(storeName); + if (storeQueue != null && !storeQueue.isEmpty()) { + internalQueuesEmptied = false; + } } } catch (ExecutionException | CancellationException e) { internalQueuesEmptied = false; AdminErrorInfo errorInfo = new AdminErrorInfo(); - int perStorePendingMessagesCount = storeAdminOperationsMapWithOffset.get(storeName).size(); + Queue storeQueue = storeAdminOperationsMapWithOffset.get(storeName); + int perStorePendingMessagesCount = storeQueue == null ? 0 : storeQueue.size(); pendingAdminMessagesCount += perStorePendingMessagesCount; storesWithPendingAdminMessagesCount += perStorePendingMessagesCount > 0 ? 1 : 0; if (e instanceof CancellationException) { - long lastSucceededId = lastSucceededExecutionIdMap.getOrDefault(storeName, -1L); - long newLastSucceededId = newLastSucceededExecutionIdMap.getOrDefault(storeName, -1L); + long lastSucceededId = lastSucceededExecutionIdMap.getOrDefault(storeName, UNASSIGNED_VALUE); + long newLastSucceededId = newLastSucceededExecutionIdMap.getOrDefault(storeName, UNASSIGNED_VALUE); - if (lastSucceededId == -1) { + if (lastSucceededId == UNASSIGNED_VALUE) { LOGGER.error("Could not find last successful execution ID for store {}", storeName); } @@ -572,22 +585,21 @@ private void executeMessagesAndCollectResults() throws InterruptedException { // only mark the store problematic if no progress is made and there are still message(s) in the queue. errorInfo.exception = new VeniceException( "Could not finish processing admin message for store " + storeName + " in time"); - errorInfo.offset = storeAdminOperationsMapWithOffset.get(storeName).peek().getOffset(); + errorInfo.offset = getNextOperationOffsetIfAvailable(storeName); problematicStores.put(storeName, errorInfo); LOGGER.warn(errorInfo.exception.getMessage()); } } else { errorInfo.exception = e; - errorInfo.offset = storeAdminOperationsMapWithOffset.get(storeName).peek().getOffset(); + errorInfo.offset = getNextOperationOffsetIfAvailable(storeName); problematicStores.put(storeName, errorInfo); } } catch (Throwable e) { - long errorMsgOffset = -1; - try { - errorMsgOffset = storeAdminOperationsMapWithOffset.get(storeName).peek().getOffset(); - } catch (Exception ex) { - LOGGER.error("Could not get the offset of the problematic admin message for store {}", storeName, ex); + long errorMsgOffset = getNextOperationOffsetIfAvailable(storeName); + if (errorMsgOffset == UNASSIGNED_VALUE) { + LOGGER.error("Could not get the offset of the problematic admin message for store {}", storeName); } + LOGGER.error( "Unexpected exception thrown while processing admin message for store {} at offset {}", storeName, @@ -632,6 +644,15 @@ private void executeMessagesAndCollectResults() throws InterruptedException { } } + /** + * @return the offset of the next enqueued operation for the given store name, or {@link #UNASSIGNED_VALUE} if unavailable. + */ + private long getNextOperationOffsetIfAvailable(String storeName) { + Queue storeQueue = storeAdminOperationsMapWithOffset.get(storeName); + AdminOperationWrapper nextOperation = storeQueue == null ? null : storeQueue.peek(); + return nextOperation == null ? UNASSIGNED_VALUE : nextOperation.getOffset(); + } + private void internalClose() { unSubscribe(); executorService.shutdownNow(); diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminExecutionTask.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminExecutionTask.java index 10f7a3cadb7..e6e6583a64b 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminExecutionTask.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminExecutionTask.java @@ -84,8 +84,6 @@ public class AdminExecutionTask implements Callable { private final ConcurrentHashMap lastSucceededExecutionIdMap; private final long lastPersistedExecutionId; - private final Map storeToScheduledTask; - AdminExecutionTask( Logger LOGGER, String clusterName, @@ -97,8 +95,7 @@ public class AdminExecutionTask implements Callable { ExecutionIdAccessor executionIdAccessor, boolean isParentController, AdminConsumptionStats stats, - String regionName, - Map storeToScheduledTask) { + String regionName) { this.LOGGER = LOGGER; this.clusterName = clusterName; this.storeName = storeName; @@ -110,7 +107,6 @@ public class AdminExecutionTask implements Callable { this.isParentController = isParentController; this.stats = stats; this.regionName = regionName; - this.storeToScheduledTask = storeToScheduledTask; } @Override @@ -159,12 +155,15 @@ public Void call() { LOGGER.error("Error {}", logMessage, e); } throw e; - } finally { - storeToScheduledTask.remove(storeName); } return null; } + // Package private for testing only + String getStoreName() { + return this.storeName; + } + private void processMessage(AdminOperation adminOperation) { long lastSucceededExecutionId = lastSucceededExecutionIdMap.getOrDefault(storeName, lastPersistedExecutionId); if (adminOperation.executionId <= lastSucceededExecutionId) { diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminOperationWrapper.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminOperationWrapper.java index ec09a8537a1..901928908cc 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminOperationWrapper.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/kafka/consumer/AdminOperationWrapper.java @@ -4,11 +4,11 @@ public class AdminOperationWrapper { - private AdminOperation adminOperation; - private long offset; - private long producerTimestamp; - private long localBrokerTimestamp; - private long delegateTimestamp; + private final AdminOperation adminOperation; + private final long offset; + private final long producerTimestamp; + private final long localBrokerTimestamp; + private final long delegateTimestamp; private Long startProcessingTimestamp = null; diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/AdminSparkServer.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/AdminSparkServer.java index c77d7fafab5..25b73ee0ac4 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/AdminSparkServer.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/AdminSparkServer.java @@ -170,6 +170,7 @@ public class AdminSparkServer extends AbstractVeniceService { private final boolean disableParentRequestTopicForStreamPushes; private final PubSubTopicRepository pubSubTopicRepository; + private final VeniceControllerRequestHandler requestHandler; public AdminSparkServer( int port, @@ -183,13 +184,15 @@ public AdminSparkServer( List disabledRoutes, VeniceProperties jettyConfigOverrides, boolean disableParentRequestTopicForStreamPushes, - PubSubTopicRepository pubSubTopicRepository) { + PubSubTopicRepository pubSubTopicRepository, + VeniceControllerRequestHandler requestHandler) { this.port = port; this.enforceSSL = enforceSSL; this.sslEnabled = sslConfig.isPresent(); this.sslConfig = sslConfig; this.checkReadMethodForKafka = checkReadMethodForKafka; this.accessController = accessController; + this.requestHandler = requestHandler; // Note: admin is passed in as a reference. The expectation is the source of the admin will // close it so we don't close it in stopInner() this.admin = admin; @@ -279,7 +282,8 @@ public boolean startInner() throws Exception { }); // Build all different routes - ControllerRoutes controllerRoutes = new ControllerRoutes(sslEnabled, accessController, pubSubTopicRepository); + ControllerRoutes controllerRoutes = + new ControllerRoutes(sslEnabled, accessController, pubSubTopicRepository, requestHandler); StoresRoutes storesRoutes = new StoresRoutes(sslEnabled, accessController, pubSubTopicRepository); JobRoutes jobRoutes = new JobRoutes(sslEnabled, accessController); SkipAdminRoute skipAdminRoute = new SkipAdminRoute(sslEnabled, accessController); @@ -362,7 +366,7 @@ public boolean startInner() throws Exception { new VeniceParentControllerRegionStateHandler(admin, createVersion.addVersionAndStartIngestion(admin))); httpService.post( NEW_STORE.getPath(), - new VeniceParentControllerRegionStateHandler(admin, createStoreRoute.createStore(admin))); + new VeniceParentControllerRegionStateHandler(admin, createStoreRoute.createStore(admin, requestHandler))); httpService.get( CHECK_RESOURCE_CLEANUP_FOR_STORE_CREATION.getPath(), new VeniceParentControllerRegionStateHandler( @@ -529,7 +533,7 @@ public boolean startInner() throws Exception { httpService.get( CLUSTER_DISCOVERY.getPath(), - new VeniceParentControllerRegionStateHandler(admin, ClusterDiscovery.discoverCluster(admin))); + new VeniceParentControllerRegionStateHandler(admin, ClusterDiscovery.discoverCluster(admin, requestHandler))); httpService.get( LIST_BOOTSTRAPPING_VERSIONS.getPath(), new VeniceParentControllerRegionStateHandler(admin, versionRoute.listBootstrappingVersions(admin))); diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/ClusterDiscovery.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/ClusterDiscovery.java index 7ddd12b55ae..91171183295 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/ClusterDiscovery.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/ClusterDiscovery.java @@ -1,12 +1,13 @@ package com.linkedin.venice.controller.server; -import static com.linkedin.venice.controllerapi.ControllerApiConstants.NAME; +import static com.linkedin.venice.controllerapi.ControllerApiConstants.STORE_NAME; import static com.linkedin.venice.controllerapi.ControllerRoute.CLUSTER_DISCOVERY; import com.linkedin.venice.HttpConstants; import com.linkedin.venice.controller.Admin; import com.linkedin.venice.controllerapi.D2ServiceDiscoveryResponse; -import com.linkedin.venice.utils.Pair; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcRequest; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcResponse; import spark.Route; @@ -14,16 +15,17 @@ public class ClusterDiscovery { /** * No ACL check; any user is allowed to discover cluster */ - public static Route discoverCluster(Admin admin) { + public static Route discoverCluster(Admin admin, VeniceControllerRequestHandler requestHandler) { return (request, response) -> { D2ServiceDiscoveryResponse responseObject = new D2ServiceDiscoveryResponse(); try { AdminSparkServer.validateParams(request, CLUSTER_DISCOVERY.getParams(), admin); - responseObject.setName(request.queryParams(NAME)); - Pair clusterToD2Pair = admin.discoverCluster(responseObject.getName()); - responseObject.setCluster(clusterToD2Pair.getFirst()); - responseObject.setD2Service(clusterToD2Pair.getSecond()); - responseObject.setServerD2Service(admin.getServerD2Service(clusterToD2Pair.getFirst())); + DiscoverClusterGrpcResponse internalResponse = requestHandler.discoverCluster( + DiscoverClusterGrpcRequest.newBuilder().setStoreName(request.queryParams(STORE_NAME)).build()); + responseObject.setName(internalResponse.getStoreName()); + responseObject.setCluster(internalResponse.getClusterName()); + responseObject.setD2Service(internalResponse.getD2Service()); + responseObject.setServerD2Service(internalResponse.getServerD2Service()); } catch (Throwable e) { responseObject.setError(e); AdminSparkServer.handleError(e, request, response); diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/ControllerRequestParamValidator.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/ControllerRequestParamValidator.java new file mode 100644 index 00000000000..5d388e2907a --- /dev/null +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/ControllerRequestParamValidator.java @@ -0,0 +1,29 @@ +package com.linkedin.venice.controller.server; + +import org.apache.commons.lang.StringUtils; + + +public class ControllerRequestParamValidator { + public static void createStoreRequestValidator( + String clusterName, + String storeName, + String owner, + String keySchema, + String valueSchema) { + if (StringUtils.isBlank(clusterName)) { + throw new IllegalArgumentException("Cluster name is required for store creation"); + } + if (StringUtils.isBlank(storeName)) { + throw new IllegalArgumentException("Store name is required for store creation"); + } + if (StringUtils.isBlank(keySchema)) { + throw new IllegalArgumentException("Key schema is required for store creation"); + } + if (StringUtils.isBlank(valueSchema)) { + throw new IllegalArgumentException("Value schema is required for store creation"); + } + if (owner == null) { + throw new IllegalArgumentException("Owner is required for store creation"); + } + } +} diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/ControllerRoutes.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/ControllerRoutes.java index 5b1e1e7fad8..6f6c7aa9180 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/ControllerRoutes.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/ControllerRoutes.java @@ -24,7 +24,8 @@ import com.linkedin.venice.controllerapi.PubSubTopicConfigResponse; import com.linkedin.venice.controllerapi.StoppableNodeStatusResponse; import com.linkedin.venice.exceptions.ErrorType; -import com.linkedin.venice.meta.Instance; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcRequest; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcResponse; import com.linkedin.venice.pubsub.PubSubTopicConfiguration; import com.linkedin.venice.pubsub.PubSubTopicRepository; import com.linkedin.venice.pubsub.api.PubSubTopic; @@ -33,6 +34,7 @@ import com.linkedin.venice.utils.Utils; import java.util.List; import java.util.Optional; +import org.apache.commons.lang.StringUtils; import org.apache.http.HttpStatus; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; @@ -45,13 +47,16 @@ public class ControllerRoutes extends AbstractRoute { private static final ObjectMapper OBJECT_MAPPER = ObjectMapperFactory.getInstance(); private final PubSubTopicRepository pubSubTopicRepository; + private final VeniceControllerRequestHandler requestHandler; public ControllerRoutes( boolean sslEnabled, Optional accessController, - PubSubTopicRepository pubSubTopicRepository) { + PubSubTopicRepository pubSubTopicRepository, + VeniceControllerRequestHandler requestHandler) { super(sslEnabled, accessController); this.pubSubTopicRepository = pubSubTopicRepository; + this.requestHandler = requestHandler; } /** @@ -64,12 +69,18 @@ public Route getLeaderController(Admin admin) { try { AdminSparkServer.validateParams(request, LEADER_CONTROLLER.getParams(), admin); String cluster = request.queryParams(CLUSTER); - responseObject.setCluster(cluster); - Instance leaderController = admin.getLeaderController(cluster); - responseObject.setUrl(leaderController.getUrl(isSslEnabled())); - if (leaderController.getPort() != leaderController.getSslPort()) { - // Controller is SSL Enabled - responseObject.setSecureUrl(leaderController.getUrl(true)); + LeaderControllerGrpcResponse inernalResponse = requestHandler + .getLeaderControllerDetails(LeaderControllerGrpcRequest.newBuilder().setClusterName(cluster).build()); + responseObject.setCluster(inernalResponse.getClusterName()); + responseObject.setUrl(inernalResponse.getHttpUrl()); + if (StringUtils.isNotBlank(inernalResponse.getHttpsUrl())) { + responseObject.setSecureUrl(inernalResponse.getHttpsUrl()); + } + if (StringUtils.isNotBlank(inernalResponse.getGrpcUrl())) { + responseObject.setGrpcUrl(inernalResponse.getGrpcUrl()); + } + if (StringUtils.isNotBlank(inernalResponse.getSecureGrpcUrl())) { + responseObject.setSecureGrpcUrl(inernalResponse.getSecureGrpcUrl()); } } catch (Throwable e) { responseObject.setError(e); diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/CreateStore.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/CreateStore.java index c0ff0ff27b5..e1a89b2dbcd 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/CreateStore.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/CreateStore.java @@ -1,5 +1,6 @@ package com.linkedin.venice.controller.server; +import static com.linkedin.venice.controller.server.VeniceControllerRequestHandler.DEFAULT_STORE_OWNER; import static com.linkedin.venice.controllerapi.ControllerApiConstants.ACCESS_PERMISSION; import static com.linkedin.venice.controllerapi.ControllerApiConstants.CLUSTER; import static com.linkedin.venice.controllerapi.ControllerApiConstants.IS_SYSTEM_STORE; @@ -18,6 +19,9 @@ import com.linkedin.venice.controllerapi.AclResponse; import com.linkedin.venice.controllerapi.ControllerResponse; import com.linkedin.venice.controllerapi.NewStoreResponse; +import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcRequest; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcResponse; import java.util.Optional; import spark.Request; import spark.Route; @@ -31,7 +35,7 @@ public CreateStore(boolean sslEnabled, Optional accessC /** * @see Admin#createStore(String, String, String, String, String, boolean, Optional) */ - public Route createStore(Admin admin) { + public Route createStore(Admin admin, VeniceControllerRequestHandler requestHandler) { return new VeniceRouteHandler(NewStoreResponse.class) { @Override public void internalHandle(Request request, NewStoreResponse veniceResponse) { @@ -39,25 +43,32 @@ public void internalHandle(Request request, NewStoreResponse veniceResponse) { if (!checkIsAllowListUser(request, veniceResponse, () -> isAllowListUser(request))) { return; } + // Validate request parameters AdminSparkServer.validateParams(request, NEW_STORE.getParams(), admin); String clusterName = request.queryParams(CLUSTER); String storeName = request.queryParams(NAME); String keySchema = request.queryParams(KEY_SCHEMA); String valueSchema = request.queryParams(VALUE_SCHEMA); boolean isSystemStore = Boolean.parseBoolean(request.queryParams(IS_SYSTEM_STORE)); - String owner = AdminSparkServer.getOptionalParameterValue(request, OWNER); if (owner == null) { - owner = ""; + owner = DEFAULT_STORE_OWNER; } - String accessPerm = request.queryParams(ACCESS_PERMISSION); - Optional accessPermissions = Optional.ofNullable(accessPerm); - veniceResponse.setCluster(clusterName); - veniceResponse.setName(storeName); - veniceResponse.setOwner(owner); - admin.createStore(clusterName, storeName, owner, keySchema, valueSchema, isSystemStore, accessPermissions); + CreateStoreGrpcRequest.Builder requestBuilder = CreateStoreGrpcRequest.newBuilder() + .setClusterStoreInfo(ClusterStoreGrpcInfo.newBuilder().setClusterName(clusterName).setStoreName(storeName)) + .setKeySchema(keySchema) + .setValueSchema(valueSchema) + .setOwner(owner) + .setIsSystemStore(isSystemStore); + if (accessPerm != null) { + requestBuilder.setAccessPermission(accessPerm); + } + CreateStoreGrpcResponse internalResponse = requestHandler.createStore(requestBuilder.build()); + veniceResponse.setCluster(internalResponse.getClusterStoreInfo().getClusterName()); + veniceResponse.setName(internalResponse.getClusterStoreInfo().getStoreName()); + veniceResponse.setOwner(internalResponse.getOwner()); } }; } diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/CreateVersion.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/CreateVersion.java index 887710f6d8a..f2bbdcaa16f 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/CreateVersion.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/CreateVersion.java @@ -33,6 +33,7 @@ import com.linkedin.venice.compression.CompressionStrategy; import com.linkedin.venice.controller.Admin; import com.linkedin.venice.controllerapi.ControllerResponse; +import com.linkedin.venice.controllerapi.RequestTopicForPushRequest; import com.linkedin.venice.controllerapi.VersionCreationResponse; import com.linkedin.venice.controllerapi.VersionResponse; import com.linkedin.venice.exceptions.ErrorType; @@ -46,11 +47,13 @@ import com.linkedin.venice.meta.Version; import com.linkedin.venice.utils.Utils; import com.linkedin.venice.utils.lazy.Lazy; -import java.security.cert.X509Certificate; +import java.util.Collections; import java.util.Optional; +import java.util.Set; import org.apache.http.HttpStatus; import org.apache.logging.log4j.LogManager; import org.apache.logging.log4j.Logger; +import spark.Request; import spark.Route; @@ -72,6 +75,380 @@ public CreateVersion( this.disableParentRequestTopicForStreamPushes = disableParentRequestTopicForStreamPushes; } + protected static void extractOptionalParamsFromRequestTopicRequest( + Request httpRequest, + RequestTopicForPushRequest request, + boolean isAclEnabled) { + request.setPartitioners(httpRequest.queryParamOrDefault(PARTITIONERS, null)); + + request.setSendStartOfPush( + Utils.parseBooleanFromString(httpRequest.queryParamOrDefault(SEND_START_OF_PUSH, "false"), SEND_START_OF_PUSH)); + + request.setSorted( + Utils.parseBooleanFromString( + httpRequest.queryParamOrDefault(PUSH_IN_SORTED_ORDER, "false"), + PUSH_IN_SORTED_ORDER)); + + request.setWriteComputeEnabled( + Utils.parseBooleanFromString( + httpRequest.queryParamOrDefault(IS_WRITE_COMPUTE_ENABLED, "false"), + IS_WRITE_COMPUTE_ENABLED)); + + request.setSeparateRealTimeTopicEnabled( + Utils.parseBooleanFromString( + httpRequest.queryParamOrDefault(SEPARATE_REAL_TIME_TOPIC_ENABLED, "false"), + SEPARATE_REAL_TIME_TOPIC_ENABLED)); + + /* + * Version-level rewind time override, and it is only valid for hybrid stores. + */ + request.setRewindTimeInSecondsOverride( + Long.parseLong(httpRequest.queryParamOrDefault(REWIND_TIME_IN_SECONDS_OVERRIDE, "-1"))); + + /* + * Version level override to defer marking this new version to the serving version post push completion. + */ + request.setDeferVersionSwap( + Utils.parseBooleanFromString(httpRequest.queryParamOrDefault(DEFER_VERSION_SWAP, "false"), DEFER_VERSION_SWAP)); + + request.setTargetedRegions(httpRequest.queryParamOrDefault(TARGETED_REGIONS, null)); + + request.setRepushSourceVersion(Integer.parseInt(httpRequest.queryParamOrDefault(REPUSH_SOURCE_VERSION, "-1"))); + + request.setSourceGridFabric(httpRequest.queryParamOrDefault(SOURCE_GRID_FABRIC, null)); + + request.setCompressionDictionary(httpRequest.queryParamOrDefault(COMPRESSION_DICTIONARY, null)); + + // Retrieve certificate from request if ACL is enabled + request.setCertificateInRequest(isAclEnabled ? getCertificate(httpRequest) : null); + } + + /** + * Verifies that the partitioner class specified in the request is valid + * based on the store's partitioner configuration. + * + *

+ * If no partitioners are provided (null or empty), the validation is skipped. + * (The store's partitioner configuration is used when no partitioners are provided). + *

+ * + * @param partitionersFromRequest An optional set of partitioners provided in the request. + * @param storePartitionerConfig The store's partitioner configuration to use for validation. + * @throws VeniceException if the store's partitioner is not in the provided partitioners. + */ + private static void validatePartitionerAgainstStoreConfig( + Set partitionersFromRequest, + PartitionerConfig storePartitionerConfig) { + // Skip validation if the user didn't provide any partitioner. Partitioner from store config will be used. + if (partitionersFromRequest == null || partitionersFromRequest.isEmpty()) { + return; + } + + // Validate if the store's partitioner matches one of the provided partitioners + if (!partitionersFromRequest.contains(storePartitionerConfig.getPartitionerClass())) { + throw new VeniceException( + "Expected partitioner class " + storePartitionerConfig.getPartitionerClass() + " cannot be found."); + } + } + + protected void verifyAndConfigurePartitionerSettings( + PartitionerConfig storePartitionerConfig, + Set partitionersFromRequest, + VersionCreationResponse response) { + validatePartitionerAgainstStoreConfig(partitionersFromRequest, storePartitionerConfig); + partitionersFromRequest = partitionersFromRequest != null ? partitionersFromRequest : Collections.emptySet(); + // Get the first partitioner that matches the store partitioner + for (String partitioner: partitionersFromRequest) { + if (storePartitionerConfig.getPartitionerClass().equals(partitioner)) { + response.setPartitionerClass(partitioner); + response.setPartitionerParams(storePartitionerConfig.getPartitionerParams()); + response.setAmplificationFactor(storePartitionerConfig.getAmplificationFactor()); + return; + } + } + response.setPartitionerClass(storePartitionerConfig.getPartitionerClass()); + response.setPartitionerParams(storePartitionerConfig.getPartitionerParams()); + response.setAmplificationFactor(storePartitionerConfig.getAmplificationFactor()); + } + + protected Lazy getActiveActiveReplicationCheck( + Admin admin, + Store store, + String clusterName, + String storeName, + boolean checkCurrentVersion) { + return Lazy.of( + () -> admin.isParent() && store.isActiveActiveReplicationEnabled() + && admin.isActiveActiveReplicationEnabledInAllRegion(clusterName, storeName, checkCurrentVersion)); + } + + protected String applyConfigBasedOnReplication( + String configType, + String configValue, + String storeName, + Lazy isActiveActiveReplicationEnabledInAllRegion) { + if (configValue != null && !isActiveActiveReplicationEnabledInAllRegion.get()) { + LOGGER.info( + "Ignoring config {} : {}, as store {} is not set up for Active/Active replication in all regions", + configType, + configValue, + storeName); + return null; + } + return configValue; + } + + /** + * Configures the source fabric to align with the native replication source fabric selection. + *

+ * For incremental pushes using a real-time (RT) policy, the push job produces to the parent Kafka cluster. + * In such cases, this method ensures that the source fabric is not overridden with the native replication (NR) + * source fabric to maintain proper configuration. + */ + protected void configureSourceFabric( + Admin admin, + Version version, + Lazy isActiveActiveReplicationEnabledInAllRegions, + RequestTopicForPushRequest request, + VersionCreationResponse response) { + PushType pushType = request.getPushType(); + // Handle native replication for non-incremental push types + if (version.isNativeReplicationEnabled() && !pushType.isIncremental()) { + String childDataCenterKafkaBootstrapServer = version.getPushStreamSourceAddress(); + if (childDataCenterKafkaBootstrapServer != null) { + response.setKafkaBootstrapServers(childDataCenterKafkaBootstrapServer); + } + response.setKafkaSourceRegion(version.getNativeReplicationSourceFabric()); + } + + // Handle incremental push with override for source region + if (admin.isParent() && pushType.isIncremental()) { + overrideSourceRegionAddressForIncrementalPushJob( + admin, + response, + request.getClusterName(), + request.getStoreName(), + request.getEmergencySourceRegion(), + request.getSourceGridFabric(), + isActiveActiveReplicationEnabledInAllRegions.get(), + version.isNativeReplicationEnabled()); + LOGGER.info( + "Using source region: {} for incremental push job: {} on store: {} cluster: {}", + response.getKafkaBootstrapServers(), + request.getPushJobId(), + request.getStoreName(), + request.getClusterName()); + } + } + + protected CompressionStrategy getCompressionStrategy(Version version, String responseTopic) { + if (Version.isRealTimeTopic(responseTopic)) { + return CompressionStrategy.NO_OP; + } + return version.getCompressionStrategy(); + } + + protected String determineResponseTopic(String storeName, Version version, RequestTopicForPushRequest request) { + String responseTopic; + PushType pushType = request.getPushType(); + if (pushType == PushType.INCREMENTAL) { + // If incremental push with a dedicated real-time topic is enabled then use the separate real-time topic + if (version.isSeparateRealTimeTopicEnabled() && request.isSeparateRealTimeTopicEnabled()) { + responseTopic = Version.composeSeparateRealTimeTopic(storeName); + } else { + responseTopic = Utils.getRealTimeTopicName(version); + } + } else if (pushType == PushType.STREAM) { + responseTopic = Version.composeRealTimeTopic(storeName); + } else if (pushType == PushType.STREAM_REPROCESSING) { + responseTopic = Version.composeStreamReprocessingTopic(storeName, version.getNumber()); + } else { + responseTopic = version.kafkaTopicName(); + } + return responseTopic; + } + + protected void handleNonStreamPushType( + Admin admin, + Store store, + RequestTopicForPushRequest request, + VersionCreationResponse response, + Lazy isActiveActiveReplicationEnabledInAllRegions) { + String clusterName = request.getClusterName(); + String storeName = request.getStoreName(); + PushType pushType = request.getPushType(); + // Check if requestTopicForPush can be handled by child controllers for the given store + if (!admin.whetherEnableBatchPushFromAdmin(storeName)) { + throw new VeniceUnsupportedOperationException( + request.getPushType().name(), + "Please push data to Venice Parent Colo instead"); + } + int computedPartitionCount = admin.calculateNumberOfPartitions(clusterName, storeName); + final Version version = admin.incrementVersionIdempotent( + clusterName, + storeName, + request.getPushJobId(), + computedPartitionCount, + response.getReplicas(), + pushType, + request.isSendStartOfPush(), + request.isSorted(), + request.getCompressionDictionary(), + Optional.ofNullable(request.getSourceGridFabric()), + Optional.ofNullable(request.getCertificateInRequest()), + request.getRewindTimeInSecondsOverride(), + Optional.ofNullable(request.getEmergencySourceRegion()), + request.isDeferVersionSwap(), + request.getTargetedRegions(), + request.getRepushSourceVersion()); + + // Set the partition count + response.setPartitions(version.getPartitionCount()); + // Set the version number + response.setVersion(version.getNumber()); + // Set the response topic + response.setKafkaTopic(determineResponseTopic(storeName, version, request)); + // Set the compression strategy + response.setCompressionStrategy(getCompressionStrategy(version, response.getKafkaTopic())); + // Set the bootstrap servers + configureSourceFabric(admin, version, isActiveActiveReplicationEnabledInAllRegions, request, response); + } + + /** + * Method handle request to get a topic for pushing data to Venice with {@link PushType#STREAM} + */ + protected void handleStreamPushType( + Admin admin, + Store store, + RequestTopicForPushRequest request, + VersionCreationResponse response, + Lazy isActiveActiveReplicationEnabledInAllRegionAllVersions) { + DataReplicationPolicy dataReplicationPolicy = store.getHybridStoreConfig().getDataReplicationPolicy(); + boolean isAggregateMode = DataReplicationPolicy.AGGREGATE.equals(dataReplicationPolicy); + if (admin.isParent()) { + // Conditionally check if the controller allows for fetching this information + if (disableParentRequestTopicForStreamPushes) { + throw new VeniceException( + "Write operations to the parent region are not permitted with push type: STREAM, as this feature is currently disabled."); + } + + // Conditionally check if this store has aggregate mode enabled. If not, throw an exception (as aggregate + // mode is required to produce to parent colo) + // We check the store config instead of the version config because we want this policy to go into effect + // without needing to perform empty pushes everywhere + if (!isAggregateMode) { + if (!isActiveActiveReplicationEnabledInAllRegionAllVersions.get()) { + throw new VeniceException( + "Store is not in aggregate mode! Cannot push data to parent topic!!. Current store setup: non-aggregate mode, AA is not enabled in all regions"); + } else { + // TODO: maybe throw exception here since this mode (REGION: PARENT, PUSH: STREAM, REPLICATION: AA-ENABLED) + // doesn't seem valid anymore + LOGGER.info( + "Store: {} samza job running in Aggregate mode; Store config is in Non-Aggregate mode; " + + "AA is enabled in all regions, letting the job continue", + store.getName()); + } + } + } else { + if (isAggregateMode) { + if (!store.isActiveActiveReplicationEnabled()) { + throw new VeniceException( + "Store is in aggregate mode and AA is not enabled. Cannot push data to child topic!!"); + } else { + LOGGER.info( + "Store: {} samza job running in Non-Aggregate mode, Store config is in Aggregate mode, " + + "AA is enabled in the local region, letting the job continue", + store.getName()); + } + } + } + + Version referenceHybridVersion = admin.getReferenceVersionForStreamingWrites( + request.getClusterName(), + request.getStoreName(), + request.getPushJobId()); + if (referenceHybridVersion == null) { + LOGGER.error( + "Request to get topic for STREAM push: {} for store: {} in cluster: {} is rejected as no hybrid version found", + request.getPushJobId(), + store.getName(), + request.getClusterName()); + throw new VeniceException( + "No hybrid version found for store: " + store.getName() + " in cluster: " + request.getClusterName() + + ". Create a hybrid version before starting a stream push job."); + } + response.setPartitions(referenceHybridVersion.getPartitionCount()); + response.setCompressionStrategy(CompressionStrategy.NO_OP); + response.setKafkaTopic(Version.composeRealTimeTopic(store.getName())); + } + + /** + * This method is used to handle the request to get a topic for pushing data to Venice. + */ + void handleRequestTopicForPushing(Admin admin, RequestTopicForPushRequest request, VersionCreationResponse response) { + String clusterName = request.getClusterName(); + String storeName = request.getStoreName(); + response.setCluster(clusterName); + response.setName(storeName); + + // Check if the store exists + Store store = admin.getStore(clusterName, storeName); + if (store == null) { + throw new VeniceNoStoreException(storeName, clusterName); + } + + // Verify and configure the partitioner + verifyAndConfigurePartitionerSettings(store.getPartitionerConfig(), request.getPartitioners(), response); + + // Validate push type + validatePushType(request.getPushType(), store); + + // Create aa replication checks with lazy evaluation + Lazy isActiveActiveReplicationEnabledInAllRegions = + getActiveActiveReplicationCheck(admin, store, clusterName, storeName, false); + Lazy isActiveActiveReplicationEnabledInAllRegionAllVersions = + getActiveActiveReplicationCheck(admin, store, clusterName, storeName, true); + + // Validate source and emergency region details and update request object + String sourceGridFabric = applyConfigBasedOnReplication( + SOURCE_GRID_FABRIC, + request.getSourceGridFabric(), + storeName, + isActiveActiveReplicationEnabledInAllRegions); + String emergencySourceRegion = applyConfigBasedOnReplication( + EMERGENCY_SOURCE_REGION, + admin.getEmergencySourceRegion(clusterName).orElse(null), + storeName, + isActiveActiveReplicationEnabledInAllRegions); + + request.setSourceGridFabric(sourceGridFabric); + request.setEmergencySourceRegion(emergencySourceRegion); + LOGGER.info( + "Request to push to store: {} in cluster: {} with source grid fabric: {} and emergency source region: {}", + storeName, + clusterName, + sourceGridFabric != null ? sourceGridFabric : "N/A", + emergencySourceRegion != null ? emergencySourceRegion : "N/A"); + + // Set the store's replication factor and partition count + response.setReplicas(admin.getReplicationFactor(clusterName, storeName)); + + boolean isSSL = admin.isSSLEnabledForPush(clusterName, storeName); + response.setKafkaBootstrapServers(admin.getKafkaBootstrapServers(isSSL)); + response.setKafkaSourceRegion(admin.getRegionName()); + response.setEnableSSL(isSSL); + + PushType pushType = request.getPushType(); + if (pushType == PushType.STREAM) { + handleStreamPushType(admin, store, request, response, isActiveActiveReplicationEnabledInAllRegionAllVersions); + } else { + handleNonStreamPushType(admin, store, request, response, isActiveActiveReplicationEnabledInAllRegions); + } + + response.setDaVinciPushStatusStoreEnabled(store.isDaVinciPushStatusStoreEnabled()); + response.setAmplificationFactor(1); + } + /** * Instead of asking Venice to create a version, pushes should ask venice which topic to write into. * The logic below includes the ability to respond with an existing topic for the same push, allowing requests @@ -108,302 +485,20 @@ public Route requestTopicForPushing(Admin admin) { return AdminSparkServer.OBJECT_MAPPER.writeValueAsString(responseObject); } + // Validate the request parameters AdminSparkServer.validateParams(request, REQUEST_TOPIC.getParams(), admin); - // Query params - String clusterName = request.queryParams(CLUSTER); - String storeName = request.queryParams(NAME); - Store store = admin.getStore(clusterName, storeName); - if (store == null) { - throw new VeniceNoStoreException(storeName); - } - responseObject.setCluster(clusterName); - responseObject.setName(storeName); - responseObject.setDaVinciPushStatusStoreEnabled(store.isDaVinciPushStatusStoreEnabled()); - - // Retrieve partitioner config from the store - PartitionerConfig storePartitionerConfig = store.getPartitionerConfig(); - if (request.queryParams(PARTITIONERS) == null) { - // Request does not contain partitioner info - responseObject.setPartitionerClass(storePartitionerConfig.getPartitionerClass()); - responseObject.setAmplificationFactor(storePartitionerConfig.getAmplificationFactor()); - responseObject.setPartitionerParams(storePartitionerConfig.getPartitionerParams()); - } else { - // Retrieve provided partitioner class list from the request - boolean hasMatchedPartitioner = false; - for (String partitioner: request.queryParams(PARTITIONERS).split(",")) { - if (partitioner.equals(storePartitionerConfig.getPartitionerClass())) { - responseObject.setPartitionerClass(storePartitionerConfig.getPartitionerClass()); - responseObject.setAmplificationFactor(storePartitionerConfig.getAmplificationFactor()); - responseObject.setPartitionerParams(storePartitionerConfig.getPartitionerParams()); - hasMatchedPartitioner = true; - break; - } - } - if (!hasMatchedPartitioner) { - throw new VeniceException( - "Expected partitioner class " + storePartitionerConfig.getPartitionerClass() + " cannot be found."); - } - } - - String pushTypeString = request.queryParams(PUSH_TYPE); - PushType pushType; - try { - pushType = PushType.valueOf(pushTypeString); - } catch (RuntimeException e) { - throw new VeniceHttpException( - HttpStatus.SC_BAD_REQUEST, - pushTypeString + " is an invalid " + PUSH_TYPE, - e, - ErrorType.BAD_REQUEST); - } - validatePushType(pushType, store); - - boolean sendStartOfPush = false; - // Make this optional so that it is compatible with old version controller client - if (request.queryParams().contains(SEND_START_OF_PUSH)) { - sendStartOfPush = Utils.parseBooleanFromString(request.queryParams(SEND_START_OF_PUSH), SEND_START_OF_PUSH); - } - - int replicationFactor = admin.getReplicationFactor(clusterName, storeName); - int partitionCount = admin.calculateNumberOfPartitions(clusterName, storeName); - responseObject.setReplicas(replicationFactor); - responseObject.setPartitions(partitionCount); - - boolean isSSL = admin.isSSLEnabledForPush(clusterName, storeName); - responseObject.setKafkaBootstrapServers(admin.getKafkaBootstrapServers(isSSL)); - responseObject.setKafkaSourceRegion(admin.getRegionName()); - responseObject.setEnableSSL(isSSL); - - String pushJobId = request.queryParams(PUSH_JOB_ID); - - boolean sorted = false; // an inefficient but safe default - String sortedParam = request.queryParams(PUSH_IN_SORTED_ORDER); - if (sortedParam != null) { - sorted = Utils.parseBooleanFromString(sortedParam, PUSH_IN_SORTED_ORDER); - } - - boolean isWriteComputeEnabled = false; - String wcEnabledParam = request.queryParams(IS_WRITE_COMPUTE_ENABLED); - if (wcEnabledParam != null) { - isWriteComputeEnabled = Utils.parseBooleanFromString(wcEnabledParam, IS_WRITE_COMPUTE_ENABLED); - } - - Optional sourceGridFabric = Optional.ofNullable(request.queryParams(SOURCE_GRID_FABRIC)); - - /** - * We can't honor source grid fabric and emergency source region config untill the store is A/A enabled in all regions. This is because - * if push job start producing to a different prod region then non A/A enabled region will not have the capability to consume from that region. - * This resets this config in such cases. - */ - Lazy isActiveActiveReplicationEnabledInAllRegion = Lazy.of(() -> { - if (admin.isParent() && store.isActiveActiveReplicationEnabled()) { - return admin.isActiveActiveReplicationEnabledInAllRegion(clusterName, storeName, false); - } else { - return false; - } - }); - - Lazy isActiveActiveReplicationEnabledInAllRegionAllVersions = Lazy.of(() -> { - if (admin.isParent() && store.isActiveActiveReplicationEnabled()) { - return admin.isActiveActiveReplicationEnabledInAllRegion(clusterName, storeName, true); - } else { - return false; - } - }); - - if (sourceGridFabric.isPresent() && !isActiveActiveReplicationEnabledInAllRegion.get()) { - LOGGER.info( - "Ignoring config {} : {}, as store {} is not set up for Active/Active replication in all regions", - SOURCE_GRID_FABRIC, - sourceGridFabric.get(), - storeName); - sourceGridFabric = Optional.empty(); - } - Optional emergencySourceRegion = admin.getEmergencySourceRegion(clusterName); - if (emergencySourceRegion.isPresent() && !isActiveActiveReplicationEnabledInAllRegion.get()) { - LOGGER.info( - "Ignoring config {} : {}, as store {} is not set up for Active/Active replication in all regions", - EMERGENCY_SOURCE_REGION, - emergencySourceRegion.get(), - storeName); - } - LOGGER.info( - "requestTopicForPushing: source grid fabric: {}, emergency source region: {}", - sourceGridFabric.orElse(""), - emergencySourceRegion.orElse("")); - - /** - * Version-level rewind time override, and it is only valid for hybrid stores. - */ - Optional rewindTimeInSecondsOverrideOptional = - Optional.ofNullable(request.queryParams(REWIND_TIME_IN_SECONDS_OVERRIDE)); - long rewindTimeInSecondsOverride = -1; - if (rewindTimeInSecondsOverrideOptional.isPresent()) { - rewindTimeInSecondsOverride = Long.parseLong(rewindTimeInSecondsOverrideOptional.get()); - } - - /** - * Version level override to defer marking this new version to the serving version post push completion. - */ - boolean deferVersionSwap = Boolean.parseBoolean(request.queryParams(DEFER_VERSION_SWAP)); - - String targetedRegions = request.queryParams(TARGETED_REGIONS); - - int repushSourceVersion = Integer.parseInt(request.queryParamOrDefault(REPUSH_SOURCE_VERSION, "-1")); - - switch (pushType) { - case BATCH: - case INCREMENTAL: - case STREAM_REPROCESSING: - if (!admin.whetherEnableBatchPushFromAdmin(storeName)) { - throw new VeniceUnsupportedOperationException( - pushTypeString, - "Please push data to Venice Parent Colo instead"); - } - String dictionaryStr = request.queryParams(COMPRESSION_DICTIONARY); - - /** - * Before trying to get the version, create the RT topic in parent kafka since it's needed anyway in following cases. - * Otherwise topic existence check fails internally. - */ - if (pushType.isIncremental() && isWriteComputeEnabled) { - admin.getRealTimeTopic(clusterName, store); - } - - final Optional certInRequest = - isAclEnabled() ? Optional.of(getCertificate(request)) : Optional.empty(); - final Version version = admin.incrementVersionIdempotent( - clusterName, - storeName, - pushJobId, - partitionCount, - replicationFactor, - pushType, - sendStartOfPush, - sorted, - dictionaryStr, - sourceGridFabric, - certInRequest, - rewindTimeInSecondsOverride, - emergencySourceRegion, - deferVersionSwap, - targetedRegions, - repushSourceVersion); - - // If Version partition count different from calculated partition count use the version count as store count - // may have been updated later. - if (version.getPartitionCount() != partitionCount) { - responseObject.setPartitions(version.getPartitionCount()); - } - String responseTopic; - /** - * Override the source fabric to respect the native replication source fabric selection. - */ - boolean overrideSourceFabric = true; - boolean isTopicRT = false; - if (pushType.isStreamReprocessing()) { - responseTopic = Version.composeStreamReprocessingTopic(storeName, version.getNumber()); - } else if (pushType.isIncremental()) { - isTopicRT = true; - if (version.isSeparateRealTimeTopicEnabled() - && Boolean.parseBoolean(request.queryParamOrDefault(SEPARATE_REAL_TIME_TOPIC_ENABLED, "false"))) { - admin.getSeparateRealTimeTopic(clusterName, storeName); - responseTopic = Version.composeSeparateRealTimeTopic(storeName); - } else { - responseTopic = Utils.getRealTimeTopicName(store); - } - // disable amplificationFactor logic on real-time topic - responseObject.setAmplificationFactor(1); - - if (version.isNativeReplicationEnabled()) { - /** - * For incremental push with RT policy store the push job produces to parent corp kafka cluster. We should not override the - * source fabric in such cases with NR source fabric. - */ - overrideSourceFabric = false; - } - } else { - responseTopic = version.kafkaTopicName(); - } - - responseObject.setVersion(version.getNumber()); - responseObject.setKafkaTopic(responseTopic); - if (isTopicRT) { - // RT topic only supports NO_OP compression - responseObject.setCompressionStrategy(CompressionStrategy.NO_OP); - } else { - responseObject.setCompressionStrategy(version.getCompressionStrategy()); - } - if (version.isNativeReplicationEnabled() && overrideSourceFabric) { - String childDataCenterKafkaBootstrapServer = version.getPushStreamSourceAddress(); - if (childDataCenterKafkaBootstrapServer != null) { - responseObject.setKafkaBootstrapServers(childDataCenterKafkaBootstrapServer); - } - responseObject.setKafkaSourceRegion(version.getNativeReplicationSourceFabric()); - } - - if (pushType.isIncremental() && admin.isParent()) { - overrideSourceRegionAddressForIncrementalPushJob( - admin, - responseObject, - clusterName, - emergencySourceRegion.orElse(null), - sourceGridFabric.orElse(null), - isActiveActiveReplicationEnabledInAllRegion.get(), - version.isNativeReplicationEnabled()); - LOGGER.info( - "Incremental push job final source region address is: {}", - responseObject.getKafkaBootstrapServers()); - } - break; - case STREAM: - - if (admin.isParent()) { - - // Conditionally check if the controller allows for fetching this information - if (disableParentRequestTopicForStreamPushes) { - throw new VeniceException( - String.format( - "Parent request topic is disabled!! Cannot push data to topic in parent colo for store %s. Aborting!!", - storeName)); - } - - // Conditionally check if this store has aggregate mode enabled. If not, throw an exception (as aggregate - // mode is required to produce to parent colo) - // We check the store config instead of the version config because we want this policy to go into affect - // without needing to perform empty pushes everywhere - if (!store.getHybridStoreConfig().getDataReplicationPolicy().equals(DataReplicationPolicy.AGGREGATE)) { - if (!isActiveActiveReplicationEnabledInAllRegionAllVersions.get()) { - throw new VeniceException("Store is not in aggregate mode! Cannot push data to parent topic!!"); - } else { - LOGGER.info( - "Store: {} samza job running in Aggregate mode, Store config is in Non-Aggregate mode, " - + "AA is enabled in all regions, letting the job continue", - storeName); - } - } - } else { - if (store.getHybridStoreConfig().getDataReplicationPolicy().equals(DataReplicationPolicy.AGGREGATE)) { - if (!store.isActiveActiveReplicationEnabled()) { - throw new VeniceException("Store is in aggregate mode! Cannot push data to child topic!!"); - } else { - LOGGER.info( - "Store: {} samza job running in Non-Aggregate mode, Store config is in Aggregate mode, " - + "AA is enabled in the local region, letting the job continue", - storeName); - } - } - } - - String realTimeTopic = admin.getRealTimeTopic(clusterName, store); - responseObject.setKafkaTopic(realTimeTopic); - // disable amplificationFactor logic on real-time topic - responseObject.setAmplificationFactor(1); - break; - default: - throw new VeniceException(pushTypeString + " is an unrecognized " + PUSH_TYPE); - } + // Extract request parameters and create a RequestTopicForPushRequest object + RequestTopicForPushRequest requestTopicForPushRequest = new RequestTopicForPushRequest( + request.queryParams(CLUSTER), + request.queryParams(NAME), + PushType.extractPushType(request.queryParams(PUSH_TYPE)), + request.queryParams(PUSH_JOB_ID)); + + // populate the request object with optional parameters + extractOptionalParamsFromRequestTopicRequest(request, requestTopicForPushRequest, isAclEnabled()); + // Invoke the handler to get the topic for pushing data + handleRequestTopicForPushing(admin, requestTopicForPushRequest, responseObject); } catch (Throwable e) { responseObject.setError(e); AdminSparkServer.handleError(e, request, response); @@ -428,6 +523,7 @@ static void overrideSourceRegionAddressForIncrementalPushJob( Admin admin, VersionCreationResponse response, String clusterName, + String storeName, String emergencySourceRegion, String pushJobSourceGridFabric, boolean isAAEnabledInAllRegions, @@ -435,7 +531,17 @@ static void overrideSourceRegionAddressForIncrementalPushJob( if (!isAAEnabledInAllRegions && isNativeReplicationEnabled) { // P2: When AA is not enabled in all the regions we use aggregate RT address, if it is available, // for inc-pushes if native-replication is enabled. - admin.getAggregateRealTimeTopicSource(clusterName).ifPresent(response::setKafkaBootstrapServers); + Optional aggregateRealTimeTopicSource = admin.getAggregateRealTimeTopicSource(clusterName); + if (aggregateRealTimeTopicSource.isPresent()) { + response.setKafkaBootstrapServers(aggregateRealTimeTopicSource.get()); + LOGGER.info( + "Incremental push job source region is being overridden with: {} address: {} for store: {} in cluster: {}", + aggregateRealTimeTopicSource.get(), + response.getKafkaBootstrapServers(), + storeName, + clusterName); + } + return; } else if (!isAAEnabledInAllRegions) { // When AA is not enabled in all regions and native replication is also disabled, don't do anything. @@ -457,13 +563,15 @@ static void overrideSourceRegionAddressForIncrementalPushJob( throw new VeniceException("Failed to get the broker server URL for the source region: " + overRideSourceRegion); } LOGGER.info( - "Incremental push job source region is being overridden with: {} address: {}", + "Incremental push job source region is being overridden with: {} address: {} for store: {} in cluster: {}", overRideSourceRegion, - bootstrapServerAddress); + bootstrapServerAddress, + storeName, + clusterName); response.setKafkaBootstrapServers(bootstrapServerAddress); } - void validatePushType(PushType pushType, Store store) { + static void validatePushType(PushType pushType, Store store) { if (pushType.equals(PushType.STREAM) && !store.isHybrid()) { throw new VeniceHttpException( HttpStatus.SC_BAD_REQUEST, diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/VeniceControllerAccessManager.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/VeniceControllerAccessManager.java new file mode 100644 index 00000000000..78a799c0cfb --- /dev/null +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/VeniceControllerAccessManager.java @@ -0,0 +1,134 @@ +package com.linkedin.venice.controller.server; + +import static java.util.Objects.requireNonNull; + +import com.linkedin.venice.acl.AclException; +import com.linkedin.venice.acl.DynamicAccessController; +import com.linkedin.venice.acl.NoOpDynamicAccessController; +import com.linkedin.venice.authorization.Method; +import java.security.cert.X509Certificate; +import javax.annotation.Nonnull; +import javax.annotation.Nullable; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +public class VeniceControllerAccessManager { + private static final Logger LOGGER = LogManager.getLogger(VeniceControllerAccessManager.class); + + protected static final String UNKNOWN_USER = "USER_UNKNOWN"; + private static final String UNKNOWN_STORE = "STORE_UNKNOWN"; + private static final String UNKNOWN_PRINCIPAL = "PRINCIPAL_UNKNOWN"; + private final DynamicAccessController accessController; + + public VeniceControllerAccessManager(DynamicAccessController accessController) { + this.accessController = requireNonNull(accessController, "DynamicAccessController is required to enforce ACL"); + } + + /** + * Checks whether the user certificate in the request grants access of the specified {@link Method} + * type to the given resource. + * + * @param resourceName The name of the resource to access. + * @param x509Certificate The user's X.509 certificate (nullable). + * @param accessMethod The method of access (e.g., GET, POST). + * @param requesterHostname The hostname of the requester (optional). + * @param requesterIp The IP address of the requester (optional). + * @return true if access is granted; false otherwise. + */ + protected boolean hasAccess( + @Nonnull String resourceName, + @Nullable X509Certificate x509Certificate, + @Nonnull Method accessMethod, + @Nullable String requesterHostname, + @Nullable String requesterIp) { + + resourceName = requireNonNull(resourceName, "Resource name is required to enforce ACL"); + accessMethod = requireNonNull(accessMethod, "Access method is required to enforce ACL"); + + try { + if (accessController.hasAccess(x509Certificate, resourceName, accessMethod.name())) { + return true; + } + // Access denied + LOGGER.warn( + "Client with principal: {} hostName: {} ipAddr: {} doesn't have access to the store: {}", + getSubjectX500Principal(x509Certificate), + requesterHostname, + requesterIp, + resourceName); + } catch (AclException e) { + LOGGER.error( + "Error when checking access for client with principal: {} hostName: {} ipAddr: {} to store: {}", + getSubjectX500Principal(x509Certificate), + requesterHostname, + requesterIp, + resourceName, + e); + } + return false; + } + + private String getSubjectX500Principal(X509Certificate x509Certificate) { + if (x509Certificate != null && x509Certificate.getSubjectX500Principal() != null) { + return x509Certificate.getSubjectX500Principal().toString(); + } + return UNKNOWN_PRINCIPAL; + } + + public boolean hasWriteAccessToPubSubTopic( + String resourceName, + X509Certificate x509Certificate, + String requesterHostname, + String requesterIp) { + return hasAccess(resourceName, x509Certificate, Method.Write, requesterHostname, requesterIp); + } + + public boolean hasReadAccessToPubSubTopic( + String resourceName, + X509Certificate x509Certificate, + String requesterHostname, + String requesterIp) { + return hasAccess(resourceName, x509Certificate, Method.Read, requesterHostname, requesterIp); + } + + public boolean hasAccessToStore( + String resourceName, + X509Certificate x509Certificate, + String requesterHostname, + String requesterIp) { + return hasAccess(resourceName, x509Certificate, Method.GET, requesterHostname, requesterIp); + } + + /** + * Check whether the user is within the admin users allowlist. + */ + public boolean isAllowListUser(String resourceName, X509Certificate x509Certificate) { + if (resourceName == null) { + resourceName = UNKNOWN_STORE; + } + return accessController.isAllowlistUsers(x509Certificate, resourceName, Method.GET.name()); + } + + public String getPrincipalId(X509Certificate x509Certificate) { + try { + if (x509Certificate != null) { + return accessController.getPrincipalId(x509Certificate); + } + LOGGER.warn("Client certificate is null. Unable to extract principal Id. Returning USER_UNKNOWN"); + } catch (Exception e) { + LOGGER.error("Error when retrieving principal Id from request", e); + } + return UNKNOWN_USER; + } + + /** + * @return whether ACL check is enabled. + */ + protected boolean isAclEnabled() { + /** + * {@link accessController} will be of type {@link NoOpDynamicAccessController} if ACL is disabled. + */ + return !(accessController instanceof NoOpDynamicAccessController); + } +} diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/VeniceControllerGrpcServiceImpl.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/VeniceControllerGrpcServiceImpl.java new file mode 100644 index 00000000000..6d28322d43b --- /dev/null +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/VeniceControllerGrpcServiceImpl.java @@ -0,0 +1,123 @@ +package com.linkedin.venice.controller.server; + +import com.linkedin.venice.controllerapi.transport.GrpcRequestResponseConverter; +import com.linkedin.venice.protocols.controller.ControllerGrpcErrorType; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcRequest; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcResponse; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcRequest; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcResponse; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcRequest; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcResponse; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcServiceGrpc.VeniceControllerGrpcServiceImplBase; +import io.grpc.Status.Code; +import io.grpc.stub.StreamObserver; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * This class is a gRPC service implementation for the VeniceController public API. + */ +public class VeniceControllerGrpcServiceImpl extends VeniceControllerGrpcServiceImplBase { + private static final Logger LOGGER = LogManager.getLogger(VeniceControllerGrpcServiceImpl.class); + + private final VeniceControllerRequestHandler requestHandler; + private final VeniceControllerAccessManager accessManager; + + public VeniceControllerGrpcServiceImpl(VeniceControllerRequestHandler requestHandler) { + this.requestHandler = requestHandler; + this.accessManager = requestHandler.getControllerAccessManager(); + } + + @Override + public void getLeaderController( + LeaderControllerGrpcRequest request, + StreamObserver responseObserver) { + String clusterName = request.getClusterName(); + LOGGER.info("Received gRPC request to get leader controller for cluster: {}", clusterName); + try { + responseObserver.onNext(requestHandler.getLeaderControllerDetails(request)); + responseObserver.onCompleted(); + } catch (IllegalArgumentException e) { + LOGGER.error("Invalid argument while getting leader controller for cluster: {}", clusterName, e); + GrpcRequestResponseConverter.sendErrorResponse( + Code.INVALID_ARGUMENT, + ControllerGrpcErrorType.BAD_REQUEST, + e, + clusterName, + null, + responseObserver); + } catch (Exception e) { + LOGGER.error("Error while getting leader controller for cluster: {}", clusterName, e); + GrpcRequestResponseConverter.sendErrorResponse( + Code.INTERNAL, + ControllerGrpcErrorType.GENERAL_ERROR, + e, + clusterName, + null, + responseObserver); + } + } + + @Override + public void discoverClusterForStore( + DiscoverClusterGrpcRequest grpcRequest, + StreamObserver responseObserver) { + String storeName = grpcRequest.getStoreName(); + LOGGER.info("Received gRPC request to discover cluster for store: {}", storeName); + try { + responseObserver.onNext(requestHandler.discoverCluster(grpcRequest)); + responseObserver.onCompleted(); + } catch (IllegalArgumentException e) { + LOGGER.error("Invalid argument while discovering cluster for store: {}", storeName, e); + GrpcRequestResponseConverter.sendErrorResponse( + Code.INVALID_ARGUMENT, + ControllerGrpcErrorType.BAD_REQUEST, + e, + null, + storeName, + responseObserver); + } catch (Exception e) { + LOGGER.error("Error while discovering cluster for store: {}", storeName, e); + GrpcRequestResponseConverter.sendErrorResponse( + Code.INTERNAL, + ControllerGrpcErrorType.GENERAL_ERROR, + e, + null, + storeName, + responseObserver); + } + } + + @Override + public void createStore( + CreateStoreGrpcRequest grpcRequest, + StreamObserver responseObserver) { + String clusterName = grpcRequest.getClusterStoreInfo().getClusterName(); + String storeName = grpcRequest.getClusterStoreInfo().getStoreName(); + LOGGER.info("Received gRPC request to create store: {} in cluster: {}", storeName, clusterName); + try { + // TODO (sushantmane) : Add the ACL check for allowlist users here + responseObserver.onNext(requestHandler.createStore(grpcRequest)); + responseObserver.onCompleted(); + } catch (IllegalArgumentException e) { + LOGGER.error("Invalid argument while creating store: {} in cluster: {}", storeName, clusterName, e); + GrpcRequestResponseConverter.sendErrorResponse( + Code.INVALID_ARGUMENT, + ControllerGrpcErrorType.BAD_REQUEST, + e, + clusterName, + storeName, + responseObserver); + } catch (Exception e) { + LOGGER.error("Error while creating store: {} in cluster: {}", storeName, clusterName, e); + GrpcRequestResponseConverter.sendErrorResponse( + Code.INTERNAL, + ControllerGrpcErrorType.GENERAL_ERROR, + e, + clusterName, + storeName, + responseObserver); + } + } +} diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/VeniceControllerRequestHandler.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/VeniceControllerRequestHandler.java new file mode 100644 index 00000000000..624f959d79a --- /dev/null +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/VeniceControllerRequestHandler.java @@ -0,0 +1,140 @@ +package com.linkedin.venice.controller.server; + +import com.linkedin.venice.controller.Admin; +import com.linkedin.venice.controller.ControllerRequestHandlerDependencies; +import com.linkedin.venice.meta.Instance; +import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcRequest; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcResponse; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcRequest; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcResponse; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcRequest; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcResponse; +import com.linkedin.venice.utils.Pair; +import java.util.Optional; +import org.apache.commons.lang.StringUtils; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * The core handler for processing incoming requests in the VeniceController. + * Acts as the central entry point for handling requests received via both HTTP/REST and gRPC protocols. + * This class is responsible for managing all request handling operations for the VeniceController. + */ +public class VeniceControllerRequestHandler { + private static final Logger LOGGER = LogManager.getLogger(VeniceControllerRequestHandler.class); + public static final String DEFAULT_STORE_OWNER = ""; + private final Admin admin; + private final boolean sslEnabled; + private final VeniceControllerAccessManager accessManager; + + public VeniceControllerRequestHandler(ControllerRequestHandlerDependencies dependencies) { + this.admin = dependencies.getAdmin(); + this.sslEnabled = dependencies.isSslEnabled(); + this.accessManager = dependencies.getControllerAccessManager(); + } + + // visibility: package-private + boolean isSslEnabled() { + return sslEnabled; + } + + /** + * The response is passed as an argument to avoid creating duplicate response objects for HTTP requests + * and to simplify unit testing with gRPC. Once the transition to gRPC is complete, we can eliminate + * the need to pass the response as an argument and instead construct and return it directly within the method. + */ + public LeaderControllerGrpcResponse getLeaderControllerDetails(LeaderControllerGrpcRequest request) { + String clusterName = request.getClusterName(); + if (StringUtils.isBlank(clusterName)) { + throw new IllegalArgumentException("Cluster name is required for leader controller discovery"); + } + Instance leaderControllerInstance = admin.getLeaderController(clusterName); + String leaderControllerUrl = leaderControllerInstance.getUrl(isSslEnabled()); + String leaderControllerSecureUrl = null; + if (leaderControllerInstance.getPort() != leaderControllerInstance.getSslPort()) { + // Controller is SSL Enabled + leaderControllerSecureUrl = leaderControllerInstance.getUrl(true); + } + LeaderControllerGrpcResponse.Builder responseBuilder = + LeaderControllerGrpcResponse.newBuilder().setClusterName(clusterName); + if (leaderControllerUrl != null) { + responseBuilder.setHttpUrl(leaderControllerUrl); + } + if (leaderControllerSecureUrl != null) { + responseBuilder.setHttpsUrl(leaderControllerSecureUrl); + } + String grpcUrl = leaderControllerInstance.getGrpcUrl(); + String secureGrpcUrl = leaderControllerInstance.getGrpcSslUrl(); + if (grpcUrl != null) { + responseBuilder.setGrpcUrl(grpcUrl); + } + if (secureGrpcUrl != null) { + responseBuilder.setSecureGrpcUrl(secureGrpcUrl); + } + return responseBuilder.build(); + } + + public DiscoverClusterGrpcResponse discoverCluster(DiscoverClusterGrpcRequest request) { + String storeName = request.getStoreName(); + if (StringUtils.isBlank(storeName)) { + throw new IllegalArgumentException("Store name is required for cluster discovery"); + } + LOGGER.info("Discovering cluster for store: {}", storeName); + Pair clusterToD2Pair = admin.discoverCluster(storeName); + + DiscoverClusterGrpcResponse.Builder responseBuilder = + DiscoverClusterGrpcResponse.newBuilder().setStoreName(storeName); + if (clusterToD2Pair.getFirst() != null) { + responseBuilder.setClusterName(clusterToD2Pair.getFirst()); + } + if (clusterToD2Pair.getSecond() != null) { + responseBuilder.setD2Service(clusterToD2Pair.getSecond()); + } + String serverD2Service = admin.getServerD2Service(clusterToD2Pair.getFirst()); + if (serverD2Service != null) { + responseBuilder.setServerD2Service(serverD2Service); + } + return responseBuilder.build(); + } + + /** + * Creates a new store in the specified Venice cluster with the provided parameters. + * @param request the request object containing all necessary details for the creation of the store + */ + public CreateStoreGrpcResponse createStore(CreateStoreGrpcRequest request) { + ClusterStoreGrpcInfo clusterStoreInfo = request.getClusterStoreInfo(); + String clusterName = clusterStoreInfo.getClusterName(); + String storeName = clusterStoreInfo.getStoreName(); + String keySchema = request.getKeySchema(); + String valueSchema = request.getValueSchema(); + String owner = request.hasOwner() ? request.getOwner() : null; + if (owner == null) { + owner = DEFAULT_STORE_OWNER; + } + Optional accessPermissions = + Optional.ofNullable(request.hasAccessPermission() ? request.getAccessPermission() : null); + boolean isSystemStore = request.hasIsSystemStore() && request.getIsSystemStore(); + ControllerRequestParamValidator.createStoreRequestValidator(clusterName, storeName, owner, keySchema, valueSchema); + LOGGER.info( + "Creating store: {} in cluster: {} with owner: {} and key schema: {} and value schema: {} and isSystemStore: {} and access permissions: {}", + storeName, + clusterName, + owner, + keySchema, + valueSchema, + isSystemStore, + accessPermissions); + admin.createStore(clusterName, storeName, owner, keySchema, valueSchema, isSystemStore, accessPermissions); + CreateStoreGrpcResponse.Builder responseBuilder = + CreateStoreGrpcResponse.newBuilder().setClusterStoreInfo(clusterStoreInfo).setOwner(owner); + + LOGGER.info("Successfully created store: {} in cluster: {}", storeName, clusterName); + return responseBuilder.build(); + } + + public VeniceControllerAccessManager getControllerAccessManager() { + return accessManager; + } +} diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/grpc/ControllerGrpcSslSessionInterceptor.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/grpc/ControllerGrpcSslSessionInterceptor.java new file mode 100644 index 00000000000..c2e801d67c7 --- /dev/null +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/grpc/ControllerGrpcSslSessionInterceptor.java @@ -0,0 +1,164 @@ +package com.linkedin.venice.controller.server.grpc; + +import com.linkedin.venice.grpc.GrpcUtils; +import com.linkedin.venice.protocols.controller.ControllerGrpcErrorType; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcErrorInfo; +import io.grpc.Context; +import io.grpc.Contexts; +import io.grpc.Grpc; +import io.grpc.Metadata; +import io.grpc.ServerCall; +import io.grpc.ServerCallHandler; +import io.grpc.ServerInterceptor; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.grpc.protobuf.StatusProto; +import java.net.SocketAddress; +import java.security.cert.X509Certificate; +import javax.net.ssl.SSLSession; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * Intercepts gRPC calls to enforce SSL/TLS requirements and propagate client certificate + * and remote address details into the gRPC {@link Context}. + * + *

If the gRPC connection does not have an SSL session, the interceptor rejects the call + * with an UNAUTHENTICATED status. Otherwise, it extracts the client certificate and remote + * address, injecting them into the gRPC context for downstream processing.

+ * + *

The following attributes are injected into the {@link Context}: + *

    + *
  • {@code CLIENT_CERTIFICATE_CONTEXT_KEY}: The client's X.509 certificate.
    • + *
  • {@code CLIENT_ADDRESS_CONTEXT_KEY}: The client's remote address as a string.
    • + *
+ * + *

Errors are logged if the SSL session is missing or if certificate extraction fails.

+ */ +public class ControllerGrpcSslSessionInterceptor implements ServerInterceptor { + private static final Logger LOGGER = LogManager.getLogger(ControllerGrpcSslSessionInterceptor.class); + protected static final String UNKNOWN_REMOTE_ADDRESS = "unknown"; + + public static final Context.Key CLIENT_CERTIFICATE_CONTEXT_KEY = + Context.key("controller-client-certificate"); + public static final Context.Key CLIENT_ADDRESS_CONTEXT_KEY = Context.key("controller-client-address"); + + protected static final VeniceControllerGrpcErrorInfo NON_SSL_ERROR_INFO = VeniceControllerGrpcErrorInfo.newBuilder() + .setStatusCode(Status.UNAUTHENTICATED.getCode().value()) + .setErrorType(ControllerGrpcErrorType.CONNECTION_ERROR) + .setErrorMessage("SSL connection required") + .build(); + + protected static final StatusRuntimeException NON_SSL_CONNECTION_ERROR = StatusProto.toStatusRuntimeException( + com.google.rpc.Status.newBuilder() + .setCode(Status.UNAUTHENTICATED.getCode().value()) + .addDetails(com.google.protobuf.Any.pack(NON_SSL_ERROR_INFO)) + .build()); + + /** + * Intercepts a gRPC call to enforce SSL/TLS requirements and propagate SSL-related attributes + * into the gRPC {@link Context}. This ensures that only secure connections with valid client + * certificates proceed further in the call chain. + * + *

The method performs the following steps: + *

    + *
  • Extracts the remote address from the server call attributes.
    • + *
  • Validates the presence of an SSL session. If absent, the call is closed with an + * {@code UNAUTHENTICATED} status.
    • + *
  • Attempts to extract the client certificate from the SSL session. If extraction fails, + * the call is closed with an {@code UNAUTHENTICATED} status.
    • + *
  • Creates a new {@link Context} containing the client certificate and remote address, + * and passes it to the downstream handlers.
    • + *
+ * + * @param serverCall The gRPC server call being intercepted. + * @param metadata The metadata associated with the call, containing headers and other request data. + * @param serverCallHandler The downstream handler that processes the call if validation passes. + * @param The request type of the gRPC method. + * @param The response type of the gRPC method. + * @return A {@link ServerCall.Listener} for handling the intercepted call, or a no-op listener + * if the call is terminated early due to validation failure. + */ + @Override + public io.grpc.ServerCall.Listener interceptCall( + ServerCall serverCall, + Metadata metadata, + ServerCallHandler serverCallHandler) { + + // Extract remote address + String remoteAddressStr = getRemoteAddress(serverCall); + + // Validate SSL session + SSLSession sslSession = serverCall.getAttributes().get(Grpc.TRANSPORT_ATTR_SSL_SESSION); + if (sslSession == null) { + return closeWithSslError(serverCall); + } + + // Extract client certificate + X509Certificate clientCert = extractClientCertificate(serverCall); + if (clientCert == null) { + return closeWithSslError(serverCall); + } + + // Create a new context with SSL-related attributes + Context context = updateAndGetContext(clientCert, remoteAddressStr); + + // Proceed with the call + return Contexts.interceptCall(context, serverCall, metadata, serverCallHandler); + } + + /** + * Retrieves the remote address from the server call attributes. + * + * @param serverCall The gRPC server call. + * @return The remote address as a string, or "unknown" if not available. + */ + private String getRemoteAddress(ServerCall serverCall) { + SocketAddress remoteAddress = serverCall.getAttributes().get(Grpc.TRANSPORT_ATTR_REMOTE_ADDR); + return remoteAddress != null ? remoteAddress.toString() : UNKNOWN_REMOTE_ADDRESS; + } + + /** + * Closes the server call with an SSL error status. + * + * @param serverCall The gRPC server call to close. + * @param The request type. + * @param The response type. + * @return A no-op listener to terminate the call. + */ + private ServerCall.Listener closeWithSslError(ServerCall serverCall) { + LOGGER.debug("SSL not enabled or client certificate extraction failed"); + serverCall.close(NON_SSL_CONNECTION_ERROR.getStatus(), NON_SSL_CONNECTION_ERROR.getTrailers()); + return new ServerCall.Listener() { + }; + } + + /** + * Extracts the client certificate from the gRPC server call. + * + * @param serverCall The gRPC server call. + * @return The client certificate, or null if extraction fails. + */ + private X509Certificate extractClientCertificate(ServerCall serverCall) { + try { + return GrpcUtils.extractGrpcClientCert(serverCall); + } catch (Exception e) { + LOGGER.error("Failed to extract client certificate", e); + return null; + } + } + + /** + * Updates the gRPC context of the current scope by adding SSL-related attributes. + * + * @param clientCert The client certificate. + * @param remoteAddressStr The remote address as a string. + * @return The new context. + */ + private Context updateAndGetContext(X509Certificate clientCert, String remoteAddressStr) { + return Context.current() + .withValue(CLIENT_CERTIFICATE_CONTEXT_KEY, clientCert) + .withValue(CLIENT_ADDRESS_CONTEXT_KEY, remoteAddressStr); + } +} diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/grpc/ParentControllerRegionValidationInterceptor.java b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/grpc/ParentControllerRegionValidationInterceptor.java new file mode 100644 index 00000000000..13222dda746 --- /dev/null +++ b/services/venice-controller/src/main/java/com/linkedin/venice/controller/server/grpc/ParentControllerRegionValidationInterceptor.java @@ -0,0 +1,67 @@ +package com.linkedin.venice.controller.server.grpc; + +import static com.linkedin.venice.controller.ParentControllerRegionState.ACTIVE; +import static com.linkedin.venice.controller.server.VeniceParentControllerRegionStateHandler.ACTIVE_CHECK_FAILURE_WARN_MESSAGE_PREFIX; + +import com.linkedin.venice.controller.Admin; +import com.linkedin.venice.controller.ParentControllerRegionState; +import com.linkedin.venice.protocols.controller.ControllerGrpcErrorType; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcErrorInfo; +import io.grpc.Grpc; +import io.grpc.Metadata; +import io.grpc.ServerCall; +import io.grpc.ServerCallHandler; +import io.grpc.ServerInterceptor; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.grpc.protobuf.StatusProto; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + + +/** + * Interceptor to verify that the parent controller is active before processing requests within its region. + */ +public class ParentControllerRegionValidationInterceptor implements ServerInterceptor { + private static final Logger LOGGER = LogManager.getLogger(ParentControllerRegionValidationInterceptor.class); + private static final VeniceControllerGrpcErrorInfo.Builder ERROR_INFO_BUILDER = + VeniceControllerGrpcErrorInfo.newBuilder() + .setErrorType(ControllerGrpcErrorType.INCORRECT_CONTROLLER) + .setStatusCode(Status.FAILED_PRECONDITION.getCode().value()); + + private static final com.google.rpc.Status.Builder RPC_STATUS_BUILDER = com.google.rpc.Status.newBuilder() + .setCode(Status.FAILED_PRECONDITION.getCode().value()) + .setMessage("Parent controller is not active"); + + private final Admin admin; + + public ParentControllerRegionValidationInterceptor(Admin admin) { + this.admin = admin; + } + + @Override + public ServerCall.Listener interceptCall( + ServerCall call, + Metadata headers, + ServerCallHandler next) { + ParentControllerRegionState parentControllerRegionState = admin.getParentControllerRegionState(); + boolean isParent = admin.isParent(); + if (isParent && parentControllerRegionState != ACTIVE) { + LOGGER.debug( + "Parent controller is not active. Rejecting the request: {} from source: {}", + call.getMethodDescriptor().getFullMethodName(), + call.getAttributes().get(Grpc.TRANSPORT_ATTR_REMOTE_ADDR)); + // Retrieve the full method name + String fullMethodName = call.getMethodDescriptor().getFullMethodName(); + VeniceControllerGrpcErrorInfo errorInfo = + ERROR_INFO_BUILDER.setErrorMessage(ACTIVE_CHECK_FAILURE_WARN_MESSAGE_PREFIX + ": " + fullMethodName).build(); + // Note: On client side convert FAILED_PRECONDITION to SC_MISDIRECTED_REQUEST + com.google.rpc.Status rpcStatus = RPC_STATUS_BUILDER.addDetails(com.google.protobuf.Any.pack(errorInfo)).build(); + StatusRuntimeException exception = StatusProto.toStatusRuntimeException(rpcStatus); + call.close(exception.getStatus(), exception.getTrailers()); + return new ServerCall.Listener() { + }; + } + return next.startCall(call, headers); + } +} diff --git a/services/venice-controller/src/main/java/com/linkedin/venice/ingestion/control/RealTimeTopicSwitcher.java b/services/venice-controller/src/main/java/com/linkedin/venice/ingestion/control/RealTimeTopicSwitcher.java index ecb0605e990..877312224ff 100644 --- a/services/venice-controller/src/main/java/com/linkedin/venice/ingestion/control/RealTimeTopicSwitcher.java +++ b/services/venice-controller/src/main/java/com/linkedin/venice/ingestion/control/RealTimeTopicSwitcher.java @@ -4,6 +4,7 @@ import static com.linkedin.venice.ConfigKeys.KAFKA_REPLICATION_FACTOR; import static com.linkedin.venice.ConfigKeys.KAFKA_REPLICATION_FACTOR_RT_TOPICS; import static com.linkedin.venice.VeniceConstants.REWIND_TIME_DECIDED_BY_SERVER; +import static com.linkedin.venice.kafka.protocol.enums.ControlMessageType.TOPIC_SWITCH; import static com.linkedin.venice.pubsub.PubSubConstants.DEFAULT_KAFKA_REPLICATION_FACTOR; import com.linkedin.venice.ConfigKeys; @@ -114,9 +115,11 @@ void sendTopicSwitch( .broadcastTopicSwitch(sourceClusters, realTimeTopic.getName(), rewindStartTimestamp, Collections.emptyMap()); } LOGGER.info( - "Successfully sent TopicSwitch into '{}' instructing to switch to '{}' with a rewindStartTimestamp of {}.", + "Successfully sent {} into '{}' instructing to switch to {} at {} with a rewindStartTimestamp of {}.", + TOPIC_SWITCH, topicWhereToSendTheTopicSwitch, realTimeTopic, + remoteKafkaUrls, rewindStartTimestamp); } @@ -135,19 +138,10 @@ void ensurePreconditions( Version version = store.getVersion(Version.parseVersionFromKafkaTopicName(topicWhereToSendTheTopicSwitch.getName())); /** - * TopicReplicator is used in child fabrics to create real-time (RT) topic when a child fabric - * is ready to start buffer replay but RT topic doesn't exist. This scenario could happen for a - * hybrid store when users haven't started any Samza job yet. In this case, RT topic should be - * created with proper retention time instead of the default 5 days retention. - * - * Potential race condition: If both rewind-time update operation and buffer-replay - * start at the same time, RT topic might not be created with the expected retention time, - * which can be fixed by sending another rewind-time update command. - * - * TODO: RT topic should be created in both parent and child fabrics when the store is converted to - * hybrid (update store command handling). However, if a store is converted to hybrid when it - * doesn't have any existing version or a correct storage quota, we cannot decide the partition - * number for it. + * We create the real-time topics when creating hybrid version for the first time. This is to ensure that the + * real-time topics are created with the correct partition count. Here we'll only check retention time and update + * it if necessary. + * TODO: Remove topic creation logic from here once new code is deployed to all regions. */ createRealTimeTopicIfNeeded(store, version, srcTopicName, hybridStoreConfig.get()); if (version != null && version.isSeparateRealTimeTopicEnabled()) { @@ -189,7 +183,6 @@ void createRealTimeTopicIfNeeded( getTopicManager().updateTopicRetention(realTimeTopic, expectedRetentionTimeMs); } } - } long getRewindStartTime( @@ -301,7 +294,8 @@ public void switchToRealTimeTopic( remoteKafkaUrls.add(aggregateRealTimeSourceKafkaUrl); } LOGGER.info( - "Will send TopicSwitch into '{}' instructing to switch to '{}' with a rewindStartTimestamp of {}.", + "Will send {} into '{}' instructing to switch to '{}' with a rewindStartTimestamp of {}.", + TOPIC_SWITCH, topicWhereToSendTheTopicSwitch, realTimeTopic, rewindStartTimestamp); diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/ControllerRequestHandlerDependenciesTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/ControllerRequestHandlerDependenciesTest.java new file mode 100644 index 00000000000..491a7878439 --- /dev/null +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/ControllerRequestHandlerDependenciesTest.java @@ -0,0 +1,105 @@ +package com.linkedin.venice.controller; + +import static org.mockito.Mockito.mock; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import com.linkedin.venice.SSLConfig; +import com.linkedin.venice.acl.DynamicAccessController; +import com.linkedin.venice.acl.NoOpDynamicAccessController; +import com.linkedin.venice.controllerapi.ControllerRoute; +import com.linkedin.venice.pubsub.PubSubTopicRepository; +import com.linkedin.venice.utils.VeniceProperties; +import io.tehuti.metrics.MetricsRepository; +import java.util.Collections; +import org.testng.annotations.Test; + + +public class ControllerRequestHandlerDependenciesTest { + @Test + public void testBuilderWithAllFieldsSet() { + Admin admin = mock(Admin.class); + SSLConfig sslConfig = mock(SSLConfig.class); + DynamicAccessController accessController = mock(DynamicAccessController.class); + PubSubTopicRepository pubSubTopicRepository = mock(PubSubTopicRepository.class); + MetricsRepository metricsRepository = mock(MetricsRepository.class); + VeniceProperties veniceProperties = mock(VeniceProperties.class); + ControllerRoute route = ControllerRoute.STORE; + + ControllerRequestHandlerDependencies dependencies = + new ControllerRequestHandlerDependencies.Builder().setAdmin(admin) + .setClusters(Collections.singleton("testCluster")) + .setEnforceSSL(true) + .setSslEnabled(true) + .setCheckReadMethodForKafka(true) + .setSslConfig(sslConfig) + .setAccessController(accessController) + .setDisabledRoutes(Collections.singletonList(route)) + .setDisableParentRequestTopicForStreamPushes(true) + .setPubSubTopicRepository(pubSubTopicRepository) + .setMetricsRepository(metricsRepository) + .setVeniceProperties(veniceProperties) + .build(); + + assertEquals(dependencies.getAdmin(), admin); + assertEquals(dependencies.getClusters(), Collections.singleton("testCluster")); + assertTrue(dependencies.isEnforceSSL()); + assertTrue(dependencies.isSslEnabled()); + assertTrue(dependencies.isCheckReadMethodForKafka()); + assertEquals(dependencies.getSslConfig(), sslConfig); + assertEquals(dependencies.getAccessController(), accessController); + assertEquals(dependencies.getDisabledRoutes(), Collections.singletonList(route)); + assertTrue(dependencies.isDisableParentRequestTopicForStreamPushes()); + assertEquals(dependencies.getPubSubTopicRepository(), pubSubTopicRepository); + assertEquals(dependencies.getMetricsRepository(), metricsRepository); + assertEquals(dependencies.getVeniceProperties(), veniceProperties); + } + + @Test + public void testBuilderWithDefaultPubSubTopicRepository() { + Admin admin = mock(Admin.class); + + ControllerRequestHandlerDependencies dependencies = + new ControllerRequestHandlerDependencies.Builder().setAdmin(admin) + .setClusters(Collections.singleton("testCluster")) + .build(); + + assertNotNull(dependencies.getPubSubTopicRepository()); + } + + @Test + public void testBuilderWithMissingAdmin() { + // Expect exception when admin is missing + IllegalArgumentException exception = expectThrows( + IllegalArgumentException.class, + () -> new ControllerRequestHandlerDependencies.Builder().setClusters(Collections.singleton("testCluster")) + .build()); + + assertEquals(exception.getMessage(), "admin is mandatory dependencies for VeniceControllerRequestHandler"); + } + + @Test + public void testDefaultValues() { + Admin admin = mock(Admin.class); + + ControllerRequestHandlerDependencies dependencies = + new ControllerRequestHandlerDependencies.Builder().setAdmin(admin) + .setClusters(Collections.singleton("testCluster")) + .build(); + + assertFalse(dependencies.isEnforceSSL()); + assertFalse(dependencies.isSslEnabled()); + assertFalse(dependencies.isCheckReadMethodForKafka()); + assertNull(dependencies.getSslConfig()); + assertNotNull(dependencies.getAccessController()); + assertTrue(dependencies.getAccessController() instanceof NoOpDynamicAccessController); + assertTrue(dependencies.getDisabledRoutes().isEmpty()); + assertFalse(dependencies.isDisableParentRequestTopicForStreamPushes()); + assertNull(dependencies.getMetricsRepository()); + assertNull(dependencies.getVeniceProperties()); + } +} diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestVeniceControllerClusterConfig.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestVeniceControllerClusterConfig.java index 42071351b4d..5764eb6b163 100644 --- a/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestVeniceControllerClusterConfig.java +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestVeniceControllerClusterConfig.java @@ -139,7 +139,7 @@ public void errOnMissingNodes() { VeniceControllerClusterConfig.parseClusterMap(builder.build(), REGION_ALLOW_LIST); } - private Properties getBaseSingleRegionProperties(boolean includeMultiRegionConfig) { + protected static Properties getBaseSingleRegionProperties(boolean includeMultiRegionConfig) { Properties props = TestUtils.getPropertiesForControllerConfig(); String clusterName = props.getProperty(CLUSTER_NAME); props.put(LOCAL_REGION_NAME, "dc-0"); diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestVeniceHelixAdmin.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestVeniceHelixAdmin.java index 5ed9d3a2452..d9b4614b573 100644 --- a/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestVeniceHelixAdmin.java +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/TestVeniceHelixAdmin.java @@ -1,32 +1,66 @@ package com.linkedin.venice.controller; +import static com.linkedin.venice.meta.Version.PushType.INCREMENTAL; +import static com.linkedin.venice.meta.Version.PushType.STREAM; +import static org.mockito.ArgumentMatchers.anyBoolean; +import static org.mockito.ArgumentMatchers.anyInt; +import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; import static org.mockito.Mockito.any; import static org.mockito.Mockito.doCallRealMethod; +import static org.mockito.Mockito.doNothing; import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; +import com.linkedin.venice.common.VeniceSystemStoreType; import com.linkedin.venice.controller.stats.DisabledPartitionStats; +import com.linkedin.venice.controller.stats.VeniceAdminStats; +import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.helix.HelixExternalViewRepository; +import com.linkedin.venice.meta.DataReplicationPolicy; +import com.linkedin.venice.meta.HybridStoreConfig; import com.linkedin.venice.meta.PartitionAssignment; import com.linkedin.venice.meta.ReadWriteStoreRepository; import com.linkedin.venice.meta.Store; import com.linkedin.venice.meta.Version; +import com.linkedin.venice.meta.Version.PushType; +import com.linkedin.venice.meta.VersionStatus; +import com.linkedin.venice.pubsub.PubSubTopicRepository; +import com.linkedin.venice.pubsub.api.PubSubTopic; +import com.linkedin.venice.pubsub.manager.TopicManager; import com.linkedin.venice.pushmonitor.ExecutionStatus; import com.linkedin.venice.utils.HelixUtils; +import com.linkedin.venice.utils.Utils; +import com.linkedin.venice.utils.locks.ClusterLockManager; import java.util.ArrayList; +import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.List; import java.util.Map; +import java.util.Optional; +import java.util.Set; +import org.mockito.ArgumentCaptor; import org.mockito.InOrder; import org.testng.annotations.Test; public class TestVeniceHelixAdmin { + private static final PubSubTopicRepository PUB_SUB_TOPIC_REPOSITORY = new PubSubTopicRepository(); + @Test public void testDropResources() { VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); @@ -58,14 +92,13 @@ public void testDropResources() { veniceHelixAdmin.deleteHelixResource(clusterName, kafkaTopic); verify(veniceHelixAdmin, times(1)).enableDisabledPartition(clusterName, kafkaTopic, false); - } /** * This test verify that in function {@link VeniceHelixAdmin#setUpMetaStoreAndMayProduceSnapshot}, * meta store RT topic creation has to happen before any writings to meta store's rt topic. * As of today, topic creation and checks to make sure that RT exists are handled in function - * {@link VeniceHelixAdmin#getRealTimeTopic}. On the other hand, as {@link VeniceHelixAdmin#storeMetadataUpdate} + * {@link VeniceHelixAdmin#ensureRealTimeTopicExistsForUserSystemStores}. On the other hand, as {@link VeniceHelixAdmin#storeMetadataUpdate} * writes to the same RT topic, it should happen after the above function. The following test enforces * such order at the statement level. * @@ -73,9 +106,9 @@ public void testDropResources() { * it is okay to relax on the ordering enforcement or delete the unit test if necessary. */ @Test - public void enforceRealTimeTopicCreationBeforeWriting() { + public void enforceRealTimeTopicCreationBeforeWritingToMetaSystemStore() { VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); - doReturn("test_rt").when(veniceHelixAdmin).getRealTimeTopic(anyString(), anyString()); + doNothing().when(veniceHelixAdmin).ensureRealTimeTopicExistsForUserSystemStores(anyString(), anyString()); doCallRealMethod().when(veniceHelixAdmin).setUpMetaStoreAndMayProduceSnapshot(anyString(), anyString()); InOrder inorder = inOrder(veniceHelixAdmin); @@ -91,8 +124,9 @@ public void enforceRealTimeTopicCreationBeforeWriting() { veniceHelixAdmin.setUpMetaStoreAndMayProduceSnapshot(anyString(), anyString()); - // Enforce that getRealTimeTopic happens before storeMetadataUpdate. See the above comments for the reasons. - inorder.verify(veniceHelixAdmin).getRealTimeTopic(anyString(), anyString()); + // Enforce that ensureRealTimeTopicExistsForUserSystemStores happens before storeMetadataUpdate. See the above + // comments for the reasons. + inorder.verify(veniceHelixAdmin).ensureRealTimeTopicExistsForUserSystemStores(anyString(), anyString()); inorder.verify(veniceHelixAdmin).storeMetadataUpdate(anyString(), anyString(), any()); } @@ -124,4 +158,734 @@ public void testGetOverallPushStatus() { overallStatus = VeniceHelixAdmin.getOverallPushStatus(veniceStatus, daVinciStatus); assertEquals(overallStatus, ExecutionStatus.DVC_INGESTION_ERROR_DISK_FULL); } + + @Test + public void testIsRealTimeTopicRequired() { + VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); + Store store = mock(Store.class, RETURNS_DEEP_STUBS); + Version version = mock(Version.class); + doCallRealMethod().when(veniceHelixAdmin).isRealTimeTopicRequired(store, version); + + // Case 1: Store is not hybrid + doReturn(false).when(store).isHybrid(); + assertFalse(veniceHelixAdmin.isRealTimeTopicRequired(store, version)); + + // Case 2: Store is hybrid and version is not hybrid + doReturn(true).when(store).isHybrid(); + doReturn(false).when(version).isHybrid(); + + // Case 3: Both store and version are hybrid && controller is child + doReturn(true).when(store).isHybrid(); + doReturn(true).when(version).isHybrid(); + assertTrue(veniceHelixAdmin.isRealTimeTopicRequired(store, version)); + doReturn(false).when(veniceHelixAdmin).isParent(); + assertTrue(veniceHelixAdmin.isRealTimeTopicRequired(store, version)); + + // Case 4: Both store and version are hybrid && controller is parent && AA is enabled + doReturn(true).when(veniceHelixAdmin).isParent(); + doReturn(true).when(store).isActiveActiveReplicationEnabled(); + assertFalse(veniceHelixAdmin.isRealTimeTopicRequired(store, version)); + + // Case 5: Both store and version are hybrid && controller is parent && AA is disabled and IncPush is enabled + doReturn(false).when(store).isActiveActiveReplicationEnabled(); + doReturn(true).when(store).isIncrementalPushEnabled(); + when(store.getHybridStoreConfig().getDataReplicationPolicy()).thenReturn(DataReplicationPolicy.NON_AGGREGATE); + assertTrue(veniceHelixAdmin.isRealTimeTopicRequired(store, version)); + + // Case 6: Both store and version are hybrid && controller is parent && AA is disabled and IncPush is disabled but + // DRP is AGGREGATE + doReturn(false).when(store).isIncrementalPushEnabled(); + when(store.getHybridStoreConfig().getDataReplicationPolicy()).thenReturn(DataReplicationPolicy.AGGREGATE); + assertTrue(veniceHelixAdmin.isRealTimeTopicRequired(store, version)); + } + + @Test + public void testCreateOrUpdateRealTimeTopics() { + String clusterName = "testCluster"; + String storeName = "testStore"; + Store store = mock(Store.class, RETURNS_DEEP_STUBS); + when(store.getName()).thenReturn(storeName); + Version version = mock(Version.class); + when(version.getStoreName()).thenReturn(storeName); + + // Case 1: Only one real-time topic is required + VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); + doCallRealMethod().when(veniceHelixAdmin).createOrUpdateRealTimeTopics(eq(clusterName), eq(store), eq(version)); + when(veniceHelixAdmin.getPubSubTopicRepository()).thenReturn(PUB_SUB_TOPIC_REPOSITORY); + doNothing().when(veniceHelixAdmin) + .createOrUpdateRealTimeTopic(eq(clusterName), eq(store), eq(version), any(PubSubTopic.class)); + veniceHelixAdmin.createOrUpdateRealTimeTopics(clusterName, store, version); + // verify and capture the arguments passed to createOrUpdateRealTimeTopic + ArgumentCaptor pubSubTopicArgumentCaptor = ArgumentCaptor.forClass(PubSubTopic.class); + verify(veniceHelixAdmin, times(1)) + .createOrUpdateRealTimeTopic(eq(clusterName), eq(store), eq(version), pubSubTopicArgumentCaptor.capture()); + assertEquals(pubSubTopicArgumentCaptor.getValue().getName(), "testStore_rt"); + + // Case 2: Both regular and separate real-time topics are required + when(version.isSeparateRealTimeTopicEnabled()).thenReturn(true); + veniceHelixAdmin.createOrUpdateRealTimeTopics(clusterName, store, version); + pubSubTopicArgumentCaptor = ArgumentCaptor.forClass(PubSubTopic.class); + // verify and capture the arguments passed to createOrUpdateRealTimeTopic + verify(veniceHelixAdmin, times(3)) + .createOrUpdateRealTimeTopic(eq(clusterName), eq(store), eq(version), pubSubTopicArgumentCaptor.capture()); + Set pubSubTopics = new HashSet<>(pubSubTopicArgumentCaptor.getAllValues()); + PubSubTopic separateRealTimeTopic = PUB_SUB_TOPIC_REPOSITORY.getTopic(storeName + "_rt_sep"); + assertTrue(pubSubTopics.contains(separateRealTimeTopic)); + } + + @Test + public void testCreateOrUpdateRealTimeTopic() { + String clusterName = "testCluster"; + String storeName = "testStore"; + int partitionCount = 10; + Store store = mock(Store.class, RETURNS_DEEP_STUBS); + when(store.getName()).thenReturn(storeName); + Version version = mock(Version.class); + when(version.getStoreName()).thenReturn(storeName); + when(version.getPartitionCount()).thenReturn(partitionCount); + PubSubTopicRepository pubSubTopicRepository = new PubSubTopicRepository(); + PubSubTopic pubSubTopic = pubSubTopicRepository.getTopic(storeName + "_rt"); + TopicManager topicManager = mock(TopicManager.class); + VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); + when(veniceHelixAdmin.getTopicManager()).thenReturn(topicManager); + + // Case 1: Real-time topic already exists + doCallRealMethod().when(veniceHelixAdmin) + .createOrUpdateRealTimeTopic(eq(clusterName), eq(store), eq(version), any(PubSubTopic.class)); + when(veniceHelixAdmin.getPubSubTopicRepository()).thenReturn(pubSubTopicRepository); + when(topicManager.containsTopic(pubSubTopic)).thenReturn(true); + doNothing().when(veniceHelixAdmin) + .validateAndUpdateTopic(eq(pubSubTopic), eq(store), eq(version), eq(partitionCount), eq(topicManager)); + veniceHelixAdmin.createOrUpdateRealTimeTopic(clusterName, store, version, pubSubTopic); + verify(veniceHelixAdmin, times(1)) + .validateAndUpdateTopic(eq(pubSubTopic), eq(store), eq(version), eq(partitionCount), eq(topicManager)); + verify(topicManager, never()).createTopic( + any(PubSubTopic.class), + anyInt(), + anyInt(), + anyLong(), + anyBoolean(), + any(Optional.class), + anyBoolean()); + + // Case 2: Real-time topic does not exist + VeniceControllerClusterConfig clusterConfig = mock(VeniceControllerClusterConfig.class); + when(veniceHelixAdmin.getControllerConfig(clusterName)).thenReturn(clusterConfig); + when(topicManager.containsTopic(pubSubTopic)).thenReturn(false); + veniceHelixAdmin.createOrUpdateRealTimeTopic(clusterName, store, version, pubSubTopic); + verify(topicManager, times(1)).createTopic( + eq(pubSubTopic), + eq(partitionCount), + anyInt(), + anyLong(), + anyBoolean(), + any(Optional.class), + anyBoolean()); + } + + @Test + public void testValidateAndUpdateTopic() { + PubSubTopic realTimeTopic = PUB_SUB_TOPIC_REPOSITORY.getTopic("testStore_rt"); + Store store = mock(Store.class, RETURNS_DEEP_STUBS); + when(store.getName()).thenReturn("testStore"); + Version version = mock(Version.class); + int expectedNumOfPartitions = 10; + TopicManager topicManager = mock(TopicManager.class); + + // Case 1: Actual partition count is not equal to expected partition count + VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); + doCallRealMethod().when(veniceHelixAdmin) + .validateAndUpdateTopic( + any(PubSubTopic.class), + any(Store.class), + any(Version.class), + anyInt(), + any(TopicManager.class)); + when(version.getPartitionCount()).thenReturn(expectedNumOfPartitions); + when(topicManager.getPartitionCount(realTimeTopic)).thenReturn(expectedNumOfPartitions - 1); + Exception exception = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin + .validateAndUpdateTopic(realTimeTopic, store, version, expectedNumOfPartitions, topicManager)); + assertTrue(exception.getMessage().contains("has different partition count")); + + // Case 2: Actual partition count is equal to expected partition count + when(topicManager.getPartitionCount(realTimeTopic)).thenReturn(expectedNumOfPartitions); + when(topicManager.updateTopicRetentionWithRetries(eq(realTimeTopic), anyLong())).thenReturn(true); + veniceHelixAdmin.validateAndUpdateTopic(realTimeTopic, store, version, expectedNumOfPartitions, topicManager); + verify(topicManager, times(1)).updateTopicRetentionWithRetries(eq(realTimeTopic), anyLong()); + } + + @Test + public void testEnsureRealTimeTopicExistsForUserSystemStores() { + String clusterName = "testCluster"; + String storeName = "testStore"; + String systemStoreName = VeniceSystemStoreType.DAVINCI_PUSH_STATUS_STORE.getSystemStoreName(storeName); + int partitionCount = 10; + Store userStore = mock(Store.class, RETURNS_DEEP_STUBS); + when(userStore.getName()).thenReturn(storeName); + Version version = mock(Version.class); + when(version.getStoreName()).thenReturn(storeName); + when(version.getPartitionCount()).thenReturn(partitionCount); + when(userStore.getPartitionCount()).thenReturn(partitionCount); + TopicManager topicManager = mock(TopicManager.class); + PubSubTopicRepository pubSubTopicRepository = new PubSubTopicRepository(); + VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); + doReturn(topicManager).when(veniceHelixAdmin).getTopicManager(); + doReturn(pubSubTopicRepository).when(veniceHelixAdmin).getPubSubTopicRepository(); + + // Case 1: Store does not exist + doReturn(null).when(veniceHelixAdmin).getStore(clusterName, storeName); + doNothing().when(veniceHelixAdmin).checkControllerLeadershipFor(clusterName); + doCallRealMethod().when(veniceHelixAdmin).ensureRealTimeTopicExistsForUserSystemStores(anyString(), anyString()); + Exception notFoundException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.ensureRealTimeTopicExistsForUserSystemStores(clusterName, storeName)); + assertTrue( + notFoundException.getMessage().contains("does not exist in"), + "Actual message: " + notFoundException.getMessage()); + + // Case 2: Store exists, but it's not user system store + doReturn(userStore).when(veniceHelixAdmin).getStore(clusterName, storeName); + Exception notUserSystemStoreException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.ensureRealTimeTopicExistsForUserSystemStores(clusterName, storeName)); + assertTrue( + notUserSystemStoreException.getMessage().contains("is not a user system store"), + "Actual message: " + notUserSystemStoreException.getMessage()); + + // Case 3: Store exists, it's a user system store, but real-time topic already exists + Store systemStore = mock(Store.class, RETURNS_DEEP_STUBS); + doReturn(systemStoreName).when(systemStore).getName(); + doReturn(Collections.emptyList()).when(systemStore).getVersions(); + doReturn(systemStore).when(veniceHelixAdmin).getStore(clusterName, systemStoreName); + doReturn(true).when(topicManager).containsTopic(any(PubSubTopic.class)); + veniceHelixAdmin.ensureRealTimeTopicExistsForUserSystemStores(clusterName, systemStoreName); + verify(topicManager, times(1)).containsTopic(any(PubSubTopic.class)); + + HelixVeniceClusterResources veniceClusterResources = mock(HelixVeniceClusterResources.class); + doReturn(veniceClusterResources).when(veniceHelixAdmin).getHelixVeniceClusterResources(clusterName); + ClusterLockManager clusterLockManager = mock(ClusterLockManager.class); + when(veniceClusterResources.getClusterLockManager()).thenReturn(clusterLockManager); + + // Case 4: Store exists, it's a user system store, first check if real-time topic exists returns false but + // later RT topic was created + topicManager = mock(TopicManager.class); + doReturn(topicManager).when(veniceHelixAdmin).getTopicManager(); + doReturn(false).doReturn(true).when(topicManager).containsTopic(any(PubSubTopic.class)); + veniceHelixAdmin.ensureRealTimeTopicExistsForUserSystemStores(clusterName, systemStoreName); + verify(topicManager, times(2)).containsTopic(any(PubSubTopic.class)); + verify(topicManager, never()).createTopic( + any(PubSubTopic.class), + anyInt(), + anyInt(), + anyLong(), + anyBoolean(), + any(Optional.class), + anyBoolean()); + + // Case 5: Store exists, it's a user system store, but real-time topic does not exist and there are no versions + // and store partition count is zero + doReturn(0).when(systemStore).getPartitionCount(); + doReturn(false).when(topicManager).containsTopic(any(PubSubTopic.class)); + Exception zeroPartitionCountException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.ensureRealTimeTopicExistsForUserSystemStores(clusterName, systemStoreName)); + assertTrue( + zeroPartitionCountException.getMessage().contains("partition count set to 0"), + "Actual message: " + zeroPartitionCountException.getMessage()); + + // Case 6: Store exists, it's a user system store, but real-time topic does not exist and there are no versions + // hence create a new real-time topic should use store's partition count + + doReturn(false).when(topicManager).containsTopic(any(PubSubTopic.class)); + doReturn(null).when(systemStore).getVersion(anyInt()); + doReturn(5).when(systemStore).getPartitionCount(); + VeniceControllerClusterConfig clusterConfig = mock(VeniceControllerClusterConfig.class); + when(veniceHelixAdmin.getControllerConfig(clusterName)).thenReturn(clusterConfig); + veniceHelixAdmin.ensureRealTimeTopicExistsForUserSystemStores(clusterName, systemStoreName); + ArgumentCaptor partitionCountArgumentCaptor = ArgumentCaptor.forClass(Integer.class); + verify(topicManager, times(1)).createTopic( + any(PubSubTopic.class), + partitionCountArgumentCaptor.capture(), + anyInt(), + anyLong(), + anyBoolean(), + any(Optional.class), + anyBoolean()); + assertEquals(partitionCountArgumentCaptor.getValue().intValue(), 5); + + // Case 7: Store exists, it's a user system store, but real-time topic does not exist and there are versions + version = mock(Version.class); + topicManager = mock(TopicManager.class); + doReturn(topicManager).when(veniceHelixAdmin).getTopicManager(); + doReturn(false).when(topicManager).containsTopic(any(PubSubTopic.class)); + doReturn(version).when(systemStore).getVersion(anyInt()); + doReturn(10).when(version).getPartitionCount(); + veniceHelixAdmin.ensureRealTimeTopicExistsForUserSystemStores(clusterName, systemStoreName); + partitionCountArgumentCaptor = ArgumentCaptor.forClass(Integer.class); + verify(topicManager, times(1)).createTopic( + any(PubSubTopic.class), + partitionCountArgumentCaptor.capture(), + anyInt(), + anyLong(), + anyBoolean(), + any(Optional.class), + anyBoolean()); + assertEquals(partitionCountArgumentCaptor.getValue().intValue(), 10); + } + + @Test + public void testValidateStoreSetupForRTWrites() { + String clusterName = "testCluster"; + String storeName = "testStore"; + String pushJobId = "pushJob123"; + VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); + Store store = mock(Store.class, RETURNS_DEEP_STUBS); + HelixVeniceClusterResources helixVeniceClusterResources = mock(HelixVeniceClusterResources.class); + ReadWriteStoreRepository storeMetadataRepository = mock(ReadWriteStoreRepository.class); + + // Mock the method chain + doReturn(helixVeniceClusterResources).when(veniceHelixAdmin).getHelixVeniceClusterResources(clusterName); + doReturn(storeMetadataRepository).when(helixVeniceClusterResources).getStoreMetadataRepository(); + doReturn(store).when(storeMetadataRepository).getStore(storeName); + + doCallRealMethod().when(veniceHelixAdmin) + .validateStoreSetupForRTWrites(anyString(), anyString(), anyString(), any(PushType.class)); + + // Case 1: Store does not exist + doReturn(null).when(storeMetadataRepository).getStore(storeName); + Exception storeNotFoundException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.validateStoreSetupForRTWrites(clusterName, storeName, pushJobId, STREAM)); + assertTrue( + storeNotFoundException.getMessage().contains("does not exist"), + "Actual message: " + storeNotFoundException.getMessage()); + + // Case 2: Store exists but is not hybrid + doReturn(store).when(storeMetadataRepository).getStore(storeName); + doReturn(false).when(store).isHybrid(); + Exception nonHybridStoreException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.validateStoreSetupForRTWrites(clusterName, storeName, pushJobId, STREAM)); + assertTrue( + nonHybridStoreException.getMessage().contains("is not a hybrid store"), + "Actual message: " + nonHybridStoreException.getMessage()); + + // Case 3: Store is hybrid but pushType is INCREMENTAL and incremental push is not enabled + doReturn(true).when(store).isHybrid(); + doReturn(false).when(store).isIncrementalPushEnabled(); + Exception incrementalPushNotEnabledException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.validateStoreSetupForRTWrites(clusterName, storeName, pushJobId, INCREMENTAL)); + assertTrue( + incrementalPushNotEnabledException.getMessage().contains("is not an incremental push store"), + "Actual message: " + incrementalPushNotEnabledException.getMessage()); + verify(store, times(1)).isIncrementalPushEnabled(); + + // Case 4: Store is hybrid and pushType is INCREMENTAL with incremental push enabled + doReturn(true).when(store).isIncrementalPushEnabled(); + veniceHelixAdmin.validateStoreSetupForRTWrites(clusterName, storeName, pushJobId, INCREMENTAL); + verify(store, times(2)).isIncrementalPushEnabled(); + } + + @Test + public void testValidateTopicPresenceAndState() { + String clusterName = "testCluster"; + String storeName = "testStore"; + String pushJobId = "pushJob123"; + PubSubTopic topic = mock(PubSubTopic.class); + int partitionCount = 10; + VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); + + TopicManager topicManager = mock(TopicManager.class); + HelixVeniceClusterResources helixVeniceClusterResources = mock(HelixVeniceClusterResources.class); + VeniceAdminStats veniceAdminStats = mock(VeniceAdminStats.class); + + doReturn(topicManager).when(veniceHelixAdmin).getTopicManager(); + doReturn(helixVeniceClusterResources).when(veniceHelixAdmin).getHelixVeniceClusterResources(clusterName); + doReturn(veniceAdminStats).when(helixVeniceClusterResources).getVeniceAdminStats(); + + doCallRealMethod().when(veniceHelixAdmin) + .validateTopicPresenceAndState( + anyString(), + anyString(), + anyString(), + any(PushType.class), + any(PubSubTopic.class), + anyInt()); + + // Case 1: Topic exists, all partitions are online, and topic is not truncated + when(topicManager.containsTopicAndAllPartitionsAreOnline(topic, partitionCount)).thenReturn(true); + when(veniceHelixAdmin.isTopicTruncated(topic.getName())).thenReturn(false); + veniceHelixAdmin + .validateTopicPresenceAndState(clusterName, storeName, pushJobId, PushType.BATCH, topic, partitionCount); + verify(topicManager, times(1)).containsTopicAndAllPartitionsAreOnline(topic, partitionCount); + verify(veniceHelixAdmin, times(1)).isTopicTruncated(topic.getName()); + + // Case 2: Topic does not exist or not all partitions are online + doReturn(false).when(topicManager).containsTopicAndAllPartitionsAreOnline(topic, partitionCount); + Exception topicAbsentException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin + .validateTopicPresenceAndState(clusterName, storeName, pushJobId, PushType.BATCH, topic, partitionCount)); + assertTrue( + topicAbsentException.getMessage().contains("is either absent or being truncated"), + "Actual message: " + topicAbsentException.getMessage()); + verify(veniceAdminStats, times(1)).recordUnexpectedTopicAbsenceCount(); + + // Case 3: Topic exists, all partitions are online, but topic is truncated + when(topicManager.containsTopicAndAllPartitionsAreOnline(topic, partitionCount)).thenReturn(true); + when(veniceHelixAdmin.isTopicTruncated(topic.getName())).thenReturn(true); + Exception topicTruncatedException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin + .validateTopicPresenceAndState(clusterName, storeName, pushJobId, INCREMENTAL, topic, partitionCount)); + assertTrue( + topicTruncatedException.getMessage().contains("is either absent or being truncated"), + "Actual message: " + topicTruncatedException.getMessage()); + verify(veniceAdminStats, times(2)).recordUnexpectedTopicAbsenceCount(); + + // Case 4: Validate behavior with different PushType (e.g., INCREMENTAL) + when(topicManager.containsTopicAndAllPartitionsAreOnline(topic, partitionCount)).thenReturn(true); + when(veniceHelixAdmin.isTopicTruncated(topic.getName())).thenReturn(false); + veniceHelixAdmin + .validateTopicPresenceAndState(clusterName, storeName, pushJobId, INCREMENTAL, topic, partitionCount); + verify(topicManager, times(4)).containsTopicAndAllPartitionsAreOnline(topic, partitionCount); + verify(veniceHelixAdmin, times(3)).isTopicTruncated(topic.getName()); + } + + @Test + public void testValidateTopicForIncrementalPush() { + String clusterName = "testCluster"; + String storeName = "testStore"; + String pushJobId = "pushJob123"; + int partitionCount = 10; + Store store = mock(Store.class); + Version referenceHybridVersion = mock(Version.class, RETURNS_DEEP_STUBS); + PubSubTopicRepository topicRepository = new PubSubTopicRepository(); + + doReturn(storeName).when(store).getName(); + doReturn(storeName).when(referenceHybridVersion).getStoreName(); + doReturn(partitionCount).when(referenceHybridVersion).getPartitionCount(); + PubSubTopic separateRtTopic = topicRepository.getTopic(Version.composeSeparateRealTimeTopic(storeName)); + PubSubTopic rtTopic = topicRepository.getTopic(Utils.getRealTimeTopicName(referenceHybridVersion)); + + VeniceHelixAdmin veniceHelixAdmin0 = mock(VeniceHelixAdmin.class); + doReturn(topicRepository).when(veniceHelixAdmin0).getPubSubTopicRepository(); + doCallRealMethod().when(veniceHelixAdmin0) + .validateTopicForIncrementalPush(anyString(), any(Store.class), any(Version.class), anyString()); + doNothing().when(veniceHelixAdmin0) + .validateTopicPresenceAndState( + anyString(), + anyString(), + anyString(), + any(PushType.class), + any(PubSubTopic.class), + anyInt()); + + // Case 1: Separate real-time topic is enabled, and both topics are valid + doReturn(true).when(referenceHybridVersion).isSeparateRealTimeTopicEnabled(); + + veniceHelixAdmin0.validateTopicForIncrementalPush(clusterName, store, referenceHybridVersion, pushJobId); + + verify(veniceHelixAdmin0, times(1)) + .validateTopicPresenceAndState(clusterName, storeName, pushJobId, INCREMENTAL, separateRtTopic, partitionCount); + verify(veniceHelixAdmin0, times(1)) + .validateTopicPresenceAndState(clusterName, storeName, pushJobId, INCREMENTAL, rtTopic, partitionCount); + + // Case 2: Separate real-time topic is disabled, only real-time topic is validated + VeniceHelixAdmin veniceHelixAdmin1 = mock(VeniceHelixAdmin.class); + doReturn(topicRepository).when(veniceHelixAdmin1).getPubSubTopicRepository(); + doCallRealMethod().when(veniceHelixAdmin1) + .validateTopicForIncrementalPush(anyString(), any(Store.class), any(Version.class), anyString()); + doNothing().when(veniceHelixAdmin1) + .validateTopicPresenceAndState( + anyString(), + anyString(), + anyString(), + any(PushType.class), + any(PubSubTopic.class), + anyInt()); + + doReturn(false).when(referenceHybridVersion).isSeparateRealTimeTopicEnabled(); + veniceHelixAdmin1.validateTopicForIncrementalPush(clusterName, store, referenceHybridVersion, pushJobId); + verify(veniceHelixAdmin1, never()) + .validateTopicPresenceAndState(clusterName, storeName, pushJobId, INCREMENTAL, separateRtTopic, partitionCount); + verify(veniceHelixAdmin1, times(1)) + .validateTopicPresenceAndState(clusterName, storeName, pushJobId, INCREMENTAL, rtTopic, partitionCount); + + // Case 3: Exception is thrown during validation of separate real-time topic + doReturn(true).when(referenceHybridVersion).isSeparateRealTimeTopicEnabled(); + doThrow(new VeniceException("Separate real-time topic validation failed")).when(veniceHelixAdmin1) + .validateTopicPresenceAndState( + anyString(), + anyString(), + anyString(), + any(PushType.class), + eq(separateRtTopic), + anyInt()); + + Exception separateRtTopicException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin1.validateTopicForIncrementalPush(clusterName, store, referenceHybridVersion, pushJobId)); + assertTrue( + separateRtTopicException.getMessage().contains("Separate real-time topic validation failed"), + "Actual message: " + separateRtTopicException.getMessage()); + + // Case 4: Exception is thrown during validation of real-time topic + doNothing().when(veniceHelixAdmin1) + .validateTopicPresenceAndState( + anyString(), + anyString(), + anyString(), + any(PushType.class), + eq(separateRtTopic), + anyInt()); + doThrow(new VeniceException("Real-time topic validation failed")).when(veniceHelixAdmin1) + .validateTopicPresenceAndState( + anyString(), + anyString(), + anyString(), + any(PushType.class), + eq(rtTopic), + anyInt()); + + Exception rtTopicException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin1.validateTopicForIncrementalPush(clusterName, store, referenceHybridVersion, pushJobId)); + assertTrue( + rtTopicException.getMessage().contains("Real-time topic validation failed"), + "Actual message: " + rtTopicException.getMessage()); + } + + @Test + public void testGetReferenceHybridVersionForRealTimeWrites() { + String clusterName = "testCluster"; + String storeName = "testStore"; + String pushJobId = "pushJob123"; + Store store = mock(Store.class, RETURNS_DEEP_STUBS); + VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); + + // Mock method calls + doReturn(storeName).when(store).getName(); + doCallRealMethod().when(veniceHelixAdmin) + .getReferenceHybridVersionForRealTimeWrites(anyString(), any(Store.class), anyString()); + + // Case 1: Store has no versions + doReturn(Collections.emptyList()).when(store).getVersions(); + Exception noVersionsException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.getReferenceHybridVersionForRealTimeWrites(clusterName, store, pushJobId)); + assertTrue( + noVersionsException.getMessage().contains("is not initialized with a version yet."), + "Actual message: " + noVersionsException.getMessage()); + + // Case 2: Store has versions, but none are valid hybrid versions + Version version1 = mock(Version.class); + Version version2 = mock(Version.class); + doReturn(Arrays.asList(version1, version2)).when(store).getVersions(); + doReturn(null).when(version1).getHybridStoreConfig(); + doReturn(null).when(version2).getHybridStoreConfig(); + + Exception noValidHybridException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.getReferenceHybridVersionForRealTimeWrites(clusterName, store, pushJobId)); + assertTrue( + noValidHybridException.getMessage().contains("No valid hybrid store version"), + "Actual message: " + noValidHybridException.getMessage()); + + // Case 3: Store has valid hybrid versions, selects the highest version number + Version validVersion1 = mock(Version.class, RETURNS_DEEP_STUBS); + Version validVersion2 = mock(Version.class, RETURNS_DEEP_STUBS); + + doReturn(10).when(validVersion1).getNumber(); + doReturn(20).when(validVersion2).getNumber(); + doReturn(VersionStatus.ONLINE).when(validVersion1).getStatus(); + doReturn(VersionStatus.ONLINE).when(validVersion2).getStatus(); + HybridStoreConfig hybridStoreConfig = mock(HybridStoreConfig.class); + doReturn(hybridStoreConfig).when(validVersion1).getHybridStoreConfig(); + doReturn(hybridStoreConfig).when(validVersion2).getHybridStoreConfig(); + doReturn(Arrays.asList(validVersion1, validVersion2)).when(store).getVersions(); + + Version referenceVersion = + veniceHelixAdmin.getReferenceHybridVersionForRealTimeWrites(clusterName, store, pushJobId); + assertEquals( + referenceVersion, + validVersion2, + "Expected the version with the highest version number to be selected."); + + // Case 4: Store has valid hybrid versions, but they are ERROR or KILLED + doReturn(VersionStatus.ERROR).when(validVersion1).getStatus(); + doReturn(VersionStatus.KILLED).when(validVersion2).getStatus(); + Exception invalidHybridVersionException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.getReferenceHybridVersionForRealTimeWrites(clusterName, store, pushJobId)); + assertTrue( + invalidHybridVersionException.getMessage().contains("No valid hybrid store version"), + "Actual message: " + invalidHybridVersionException.getMessage()); + } + + @Test + public void testGetIncrementalPushVersion() { + String clusterName = "testCluster"; + String storeName = "testStore"; + String pushJobId = "pushJob123"; + + VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); + HelixVeniceClusterResources resources = mock(HelixVeniceClusterResources.class, RETURNS_DEEP_STUBS); + ClusterLockManager lockManager = new ClusterLockManager(clusterName); + Store store = mock(Store.class); + Version hybridVersion = mock(Version.class); + + doReturn(resources).when(veniceHelixAdmin).getHelixVeniceClusterResources(clusterName); + doReturn(lockManager).when(resources).getClusterLockManager(); + doReturn(hybridVersion).when(veniceHelixAdmin) + .getReferenceHybridVersionForRealTimeWrites(clusterName, store, pushJobId); + + when(resources.getStoreMetadataRepository().getStore(storeName)).thenReturn(store); + + doCallRealMethod().when(veniceHelixAdmin).getIncrementalPushVersion(anyString(), anyString(), anyString()); + + // Case 1: All validations pass, and the real-time topic is not required + doNothing().when(veniceHelixAdmin).checkControllerLeadershipFor(clusterName); + doNothing().when(veniceHelixAdmin) + .validateStoreSetupForRTWrites(eq(clusterName), eq(storeName), eq(pushJobId), eq(INCREMENTAL)); + doReturn(false).when(veniceHelixAdmin).isRealTimeTopicRequired(store, hybridVersion); + Version result = veniceHelixAdmin.getIncrementalPushVersion(clusterName, storeName, pushJobId); + assertEquals(result, hybridVersion, "Expected the hybrid version to be returned."); + + // Case 2: Real-time topic is required, and validation succeeds + doReturn(true).when(veniceHelixAdmin).isRealTimeTopicRequired(store, hybridVersion); + doNothing().when(veniceHelixAdmin).validateTopicForIncrementalPush(clusterName, store, hybridVersion, pushJobId); + result = veniceHelixAdmin.getIncrementalPushVersion(clusterName, storeName, pushJobId); + assertEquals(result, hybridVersion, "Expected the hybrid version to be returned after topic validation."); + verify(veniceHelixAdmin, times(1)).validateTopicForIncrementalPush(clusterName, store, hybridVersion, pushJobId); + + // Case 3: Real-time topic validation fails + doThrow(new VeniceException("Real-time topic validation failed")).when(veniceHelixAdmin) + .validateTopicForIncrementalPush(clusterName, store, hybridVersion, pushJobId); + Exception rtTopicValidationException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.getIncrementalPushVersion(clusterName, storeName, pushJobId)); + assertTrue( + rtTopicValidationException.getMessage().contains("Real-time topic validation failed"), + "Actual message: " + rtTopicValidationException.getMessage()); + + // Case 4: Reference hybrid version retrieval fails + doThrow(new VeniceException("No valid hybrid version found")).when(veniceHelixAdmin) + .getReferenceHybridVersionForRealTimeWrites(clusterName, store, pushJobId); + Exception hybridVersionException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.getIncrementalPushVersion(clusterName, storeName, pushJobId)); + assertTrue( + hybridVersionException.getMessage().contains("No valid hybrid version found"), + "Actual message: " + hybridVersionException.getMessage()); + + // Case 5: Store setup validation fails + doThrow(new VeniceException("Store setup validation failed")).when(veniceHelixAdmin) + .validateStoreSetupForRTWrites(eq(clusterName), eq(storeName), eq(pushJobId), eq(INCREMENTAL)); + + Exception storeSetupValidationException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.getIncrementalPushVersion(clusterName, storeName, pushJobId)); + assertTrue( + storeSetupValidationException.getMessage().contains("Store setup validation failed"), + "Actual message: " + storeSetupValidationException.getMessage()); + } + + @Test + public void testGetReferenceVersionForStreamingWrites() { + String clusterName = "testCluster"; + String storeName = "testStore"; + String pushJobId = "pushJob123"; + int partitionCount = 10; + + VeniceHelixAdmin veniceHelixAdmin = mock(VeniceHelixAdmin.class); + HelixVeniceClusterResources resources = mock(HelixVeniceClusterResources.class, RETURNS_DEEP_STUBS); + ClusterLockManager lockManager = new ClusterLockManager(clusterName); + Store store = mock(Store.class); + Version hybridVersion = mock(Version.class); + PubSubTopicRepository topicRepository = new PubSubTopicRepository(); + PubSubTopic rtTopic = topicRepository.getTopic("testStore_rt"); + + doReturn(storeName).when(store).getName(); + doReturn(storeName).when(hybridVersion).getStoreName(); + doReturn(resources).when(veniceHelixAdmin).getHelixVeniceClusterResources(clusterName); + doReturn(lockManager).when(resources).getClusterLockManager(); + doReturn(hybridVersion).when(veniceHelixAdmin) + .getReferenceHybridVersionForRealTimeWrites(clusterName, store, pushJobId); + doReturn(topicRepository).when(veniceHelixAdmin).getPubSubTopicRepository(); + doReturn(partitionCount).when(hybridVersion).getPartitionCount(); + + when(resources.getStoreMetadataRepository().getStore(storeName)).thenReturn(store); + + doCallRealMethod().when(veniceHelixAdmin) + .getReferenceVersionForStreamingWrites(anyString(), anyString(), anyString()); + + // Case 1: All validations pass, and the controller is parent + doNothing().when(veniceHelixAdmin).checkControllerLeadershipFor(clusterName); + doNothing().when(veniceHelixAdmin) + .validateStoreSetupForRTWrites(eq(clusterName), eq(storeName), eq(pushJobId), eq(PushType.STREAM)); + doReturn(true).when(veniceHelixAdmin).isParent(); + + Version result = veniceHelixAdmin.getReferenceVersionForStreamingWrites(clusterName, storeName, pushJobId); + assertEquals(result, hybridVersion, "Expected the hybrid version to be returned."); + + // Case 2: All validations pass, and the controller is not parent + doReturn(false).when(veniceHelixAdmin).isParent(); + doNothing().when(veniceHelixAdmin) + .validateTopicPresenceAndState( + eq(clusterName), + eq(storeName), + eq(pushJobId), + eq(PushType.STREAM), + eq(rtTopic), + eq(partitionCount)); + result = veniceHelixAdmin.getReferenceVersionForStreamingWrites(clusterName, storeName, pushJobId); + assertEquals(result, hybridVersion, "Expected the hybrid version to be returned after topic validation."); + verify(veniceHelixAdmin, times(1)).validateTopicPresenceAndState( + eq(clusterName), + eq(storeName), + eq(pushJobId), + eq(PushType.STREAM), + eq(rtTopic), + eq(partitionCount)); + + // Case 3: Topic validation fails when the controller is not parent + doThrow(new VeniceException("Real-time topic validation failed")).when(veniceHelixAdmin) + .validateTopicPresenceAndState( + eq(clusterName), + eq(storeName), + eq(pushJobId), + eq(PushType.STREAM), + eq(rtTopic), + eq(partitionCount)); + Exception rtTopicValidationException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.getReferenceVersionForStreamingWrites(clusterName, storeName, pushJobId)); + assertTrue( + rtTopicValidationException.getMessage().contains("Real-time topic validation failed"), + "Actual message: " + rtTopicValidationException.getMessage()); + + // Case 4: Reference hybrid version retrieval fails + doThrow(new VeniceException("No valid hybrid version found")).when(veniceHelixAdmin) + .getReferenceHybridVersionForRealTimeWrites(clusterName, store, pushJobId); + Exception hybridVersionException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.getReferenceVersionForStreamingWrites(clusterName, storeName, pushJobId)); + assertTrue( + hybridVersionException.getMessage().contains("No valid hybrid version found"), + "Actual message: " + hybridVersionException.getMessage()); + + // Case 5: Store setup validation fails + doThrow(new VeniceException("Store setup validation failed")).when(veniceHelixAdmin) + .validateStoreSetupForRTWrites(eq(clusterName), eq(storeName), eq(pushJobId), eq(PushType.STREAM)); + Exception storeSetupValidationException = expectThrows( + VeniceException.class, + () -> veniceHelixAdmin.getReferenceVersionForStreamingWrites(clusterName, storeName, pushJobId)); + assertTrue( + storeSetupValidationException.getMessage().contains("Store setup validation failed"), + "Actual message: " + storeSetupValidationException.getMessage()); + } } diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTaskTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTaskTest.java index be4623c3ba8..71fcaa5034f 100644 --- a/services/venice-controller/src/test/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTaskTest.java +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/kafka/consumer/AdminConsumptionTaskTest.java @@ -60,7 +60,6 @@ import com.linkedin.venice.controller.stats.AdminConsumptionStats; import com.linkedin.venice.exceptions.VeniceException; import com.linkedin.venice.guid.GuidUtils; -import com.linkedin.venice.kafka.protocol.KafkaMessageEnvelope; import com.linkedin.venice.kafka.protocol.state.PartitionState; import com.linkedin.venice.kafka.protocol.state.ProducerPartitionState; import com.linkedin.venice.kafka.validation.SegmentStatus; @@ -73,7 +72,6 @@ import com.linkedin.venice.pubsub.PubSubTopicPartitionImpl; import com.linkedin.venice.pubsub.PubSubTopicRepository; import com.linkedin.venice.pubsub.api.PubSubConsumerAdapter; -import com.linkedin.venice.pubsub.api.PubSubMessageDeserializer; import com.linkedin.venice.pubsub.api.PubSubProduceResult; import com.linkedin.venice.pubsub.api.PubSubTopic; import com.linkedin.venice.pubsub.api.PubSubTopicPartition; @@ -81,7 +79,6 @@ import com.linkedin.venice.pubsub.manager.TopicManager; import com.linkedin.venice.serialization.DefaultSerializer; import com.linkedin.venice.serialization.avro.AvroProtocolDefinition; -import com.linkedin.venice.serialization.avro.OptimizedKafkaValueSerializer; import com.linkedin.venice.unit.kafka.InMemoryKafkaBroker; import com.linkedin.venice.unit.kafka.SimplePartitioner; import com.linkedin.venice.unit.kafka.consumer.MockInMemoryConsumer; @@ -100,22 +97,25 @@ import com.linkedin.venice.utils.Utils; import com.linkedin.venice.utils.VeniceProperties; import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap; -import com.linkedin.venice.utils.pools.LandFillObjectPool; import com.linkedin.venice.writer.VeniceWriter; import com.linkedin.venice.writer.VeniceWriterOptions; import java.io.IOException; import java.nio.ByteBuffer; import java.util.ArrayList; import java.util.Arrays; +import java.util.Collection; import java.util.Collections; 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.Properties; import java.util.Queue; import java.util.Set; +import java.util.concurrent.Callable; +import java.util.concurrent.CancellationException; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; @@ -123,7 +123,10 @@ import java.util.concurrent.TimeUnit; import java.util.concurrent.TimeoutException; import org.mockito.AdditionalAnswers; +import org.mockito.ArgumentCaptor; import org.mockito.Mockito; +import org.mockito.internal.verification.VerificationModeFactory; +import org.mockito.verification.Timeout; import org.testng.Assert; import org.testng.annotations.AfterMethod; import org.testng.annotations.BeforeMethod; @@ -235,10 +238,6 @@ private AdminConsumptionTask getAdminConsumptionTask( new MockInMemoryConsumer(inMemoryKafkaBroker, pollStrategy, mockKafkaConsumer); PubSubTopicRepository pubSubTopicRepository = new PubSubTopicRepository(); - PubSubMessageDeserializer pubSubMessageDeserializer = new PubSubMessageDeserializer( - new OptimizedKafkaValueSerializer(), - new LandFillObjectPool<>(KafkaMessageEnvelope::new), - new LandFillObjectPool<>(KafkaMessageEnvelope::new)); return new AdminConsumptionTask( clusterName, @@ -255,7 +254,6 @@ private AdminConsumptionTask getAdminConsumptionTask( adminConsumptionCycleTimeoutMs, maxWorkerThreadPoolSize, pubSubTopicRepository, - pubSubMessageDeserializer, "dc-0"); } @@ -1757,4 +1755,43 @@ public void testSystemStoreMessageOrder() throws InterruptedException, IOExcepti executor.shutdown(); executor.awaitTermination(TIMEOUT, TimeUnit.MILLISECONDS); } + + /** + * In some edge cases the AdminExecutionTask may be slow or stuck and eventually timeout. We want to ensure the + * unprocessed admin messages are queued up again during the next cycle. + */ + @Test(timeOut = TIMEOUT) + public void testTimedOutAdminExecutionTask() throws IOException, InterruptedException, ExecutionException { + ExecutorService mockExecutorService = mock(ExecutorService.class); + AdminConsumptionStats mockStats = mock(AdminConsumptionStats.class); + List> results = new ArrayList<>(); + Future mockFuture = mock(Future.class); + doThrow(new CancellationException("timed out")).when(mockFuture).get(); + results.add(mockFuture); + doReturn(results).when(mockExecutorService).invokeAll(any(), anyLong(), any()); + AdminConsumptionTask task = getAdminConsumptionTask(new RandomPollStrategy(), false, mockStats, 100); + task.setAdminExecutionTaskExecutorService(mockExecutorService); + veniceWriter.put( + emptyKeyBytes, + getStoreCreationMessage(clusterName, storeName, owner, keySchema, valueSchema, 1L), + AdminOperationSerializer.LATEST_SCHEMA_ID_FOR_ADMIN_OPERATION); + executor.submit(task); + ArgumentCaptor>> tasksCaptor = ArgumentCaptor.forClass(Collection.class); + verify(mockExecutorService, new Timeout(5000, VerificationModeFactory.times(2))) + .invokeAll(tasksCaptor.capture(), anyLong(), any()); + // The timed out or cancelled execution task should still be queued up again and again + for (Collection> tasks: tasksCaptor.getAllValues()) { + Assert.assertEquals(tasks.size(), 1); + Assert.assertTrue(tasks.iterator().next() instanceof AdminExecutionTask); + AdminExecutionTask adminExecutionTask = (AdminExecutionTask) tasks.iterator().next(); + Assert.assertEquals(adminExecutionTask.getStoreName(), storeName); + } + verify(mockStats, atLeastOnce()).recordPendingAdminMessagesCount(1d); + verify(mockStats, atLeastOnce()).recordStoresWithPendingAdminMessagesCount(1d); + Assert.assertNotNull(task.getLastExceptionForStore(storeName)); + Assert.assertTrue(task.getLastExceptionForStore(storeName).getMessage().contains("Could not finish processing")); + task.close(); + executor.shutdown(); + executor.awaitTermination(TIMEOUT, TimeUnit.MILLISECONDS); + } } diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/ClusterDiscoveryTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/ClusterDiscoveryTest.java new file mode 100644 index 00000000000..7448e854aa1 --- /dev/null +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/ClusterDiscoveryTest.java @@ -0,0 +1,69 @@ +package com.linkedin.venice.controller.server; + +import static com.linkedin.venice.controllerapi.ControllerRoute.CLUSTER_DISCOVERY; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + +import com.linkedin.venice.controller.Admin; +import com.linkedin.venice.controllerapi.ControllerApiConstants; +import com.linkedin.venice.controllerapi.D2ServiceDiscoveryResponse; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcRequest; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcResponse; +import com.linkedin.venice.utils.ObjectMapperFactory; +import org.testng.annotations.Test; +import spark.QueryParamsMap; +import spark.Request; +import spark.Response; +import spark.Route; + + +public class ClusterDiscoveryTest { + @Test + public void testDiscoverCluster() throws Exception { + // Case 1: Store name is not provided + String clusterName = "test-cluster"; + String storeName = "test-store"; + String d2Service = "d2://test-service"; + String serverD2Service = "d2://test-server"; + + Admin admin = mock(Admin.class); + VeniceControllerRequestHandler requestHandler = mock(VeniceControllerRequestHandler.class); + + DiscoverClusterGrpcResponse discoverClusterGrpcResponse = DiscoverClusterGrpcResponse.newBuilder() + .setClusterName(clusterName) + .setStoreName(storeName) + .setD2Service(d2Service) + .setServerD2Service(serverD2Service) + .build(); + doReturn(discoverClusterGrpcResponse).when(requestHandler).discoverCluster(any(DiscoverClusterGrpcRequest.class)); + + Request request = mock(Request.class); + when(request.pathInfo()).thenReturn(CLUSTER_DISCOVERY.getPath()); + when(request.queryParams(eq(ControllerApiConstants.NAME))).thenReturn(storeName); + Response response = mock(Response.class); + + Route discoverCluster = ClusterDiscovery.discoverCluster(admin, requestHandler); + D2ServiceDiscoveryResponse d2ServiceDiscoveryResponse = ObjectMapperFactory.getInstance() + .readValue(discoverCluster.handle(request, response).toString(), D2ServiceDiscoveryResponse.class); + assertNotNull(d2ServiceDiscoveryResponse, "Response should not be null"); + assertEquals(d2ServiceDiscoveryResponse.getName(), storeName, "Store name should match"); + assertEquals(d2ServiceDiscoveryResponse.getCluster(), clusterName, "Cluster name should match"); + assertEquals(d2ServiceDiscoveryResponse.getD2Service(), d2Service, "D2 service should match"); + assertEquals(d2ServiceDiscoveryResponse.getServerD2Service(), serverD2Service, "Server D2 service should match"); + + // Case 2: Store name is not provided + QueryParamsMap queryParamsMap = mock(QueryParamsMap.class); + when(request.queryMap()).thenReturn(queryParamsMap); + when(request.queryParams(eq(ControllerApiConstants.NAME))).thenReturn(""); + D2ServiceDiscoveryResponse d2ServiceDiscoveryResponse2 = ObjectMapperFactory.getInstance() + .readValue(discoverCluster.handle(request, response).toString(), D2ServiceDiscoveryResponse.class); + assertNotNull(d2ServiceDiscoveryResponse2, "Response should not be null"); + assertTrue(d2ServiceDiscoveryResponse2.isError(), "Error should be present in the response"); + } +} diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/ControllerRequestParamValidatorTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/ControllerRequestParamValidatorTest.java new file mode 100644 index 00000000000..ca89a0be39e --- /dev/null +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/ControllerRequestParamValidatorTest.java @@ -0,0 +1,47 @@ +package com.linkedin.venice.controller.server; + +import org.testng.annotations.DataProvider; +import org.testng.annotations.Test; + + +public class ControllerRequestParamValidatorTest { + @Test(dataProvider = "validInputProvider") + public void testCreateStoreRequestValidatorValidInputs( + String clusterName, + String storeName, + String owner, + String keySchema, + String valueSchema) { + ControllerRequestParamValidator.createStoreRequestValidator(clusterName, storeName, owner, keySchema, valueSchema); + } + + @Test(dataProvider = "missingParameterProvider", expectedExceptions = IllegalArgumentException.class) + public void testCreateStoreRequestValidatorMissingParameters( + String clusterName, + String storeName, + String owner, + String keySchema, + String valueSchema) { + ControllerRequestParamValidator.createStoreRequestValidator(clusterName, storeName, owner, keySchema, valueSchema); + } + + @DataProvider + public Object[][] validInputProvider() { + return new Object[][] { { "clusterA", "storeA", "ownerA", "keySchemaA", "valueSchemaA" }, + { "clusterB", "storeB", "ownerB", "keySchemaB", "valueSchemaB" } }; + } + + @DataProvider + public Object[][] missingParameterProvider() { + return new Object[][] { { null, "storeA", "ownerA", "keySchemaA", "valueSchemaA" }, // Missing clusterName + { "", "storeA", "ownerA", "keySchemaA", "valueSchemaA" }, // Empty clusterName + { "clusterA", null, "ownerA", "keySchemaA", "valueSchemaA" }, // Missing storeName + { "clusterA", "", "ownerA", "keySchemaA", "valueSchemaA" }, // Empty storeName + { "clusterA", "storeA", null, "keySchemaA", "valueSchemaA" }, // Missing owner + { "clusterA", "storeA", "ownerA", null, "valueSchemaA" }, // Missing keySchema + { "clusterA", "storeA", "ownerA", "", "valueSchemaA" }, // Empty keySchema + { "clusterA", "storeA", "ownerA", "keySchemaA", null }, // Missing valueSchema + { "clusterA", "storeA", "ownerA", "keySchemaA", "" } // Empty valueSchema + }; + } +} diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/ControllerRoutesTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/ControllerRoutesTest.java index 5b4b31d76e4..053e32611a6 100644 --- a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/ControllerRoutesTest.java +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/ControllerRoutesTest.java @@ -5,11 +5,13 @@ import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertTrue; import com.fasterxml.jackson.databind.ObjectMapper; import com.linkedin.venice.controller.Admin; +import com.linkedin.venice.controller.ControllerRequestHandlerDependencies; import com.linkedin.venice.controller.InstanceRemovableStatuses; import com.linkedin.venice.controller.VeniceParentHelixAdmin; import com.linkedin.venice.controllerapi.AggregatedHealthStatusRequest; @@ -24,6 +26,7 @@ import java.util.List; import java.util.Map; import java.util.Optional; +import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; import spark.Request; import spark.Response; @@ -37,39 +40,62 @@ public class ControllerRoutesTest { private static final String TEST_HOST = "localhost"; private static final int TEST_PORT = 2181; private static final int TEST_SSL_PORT = 2182; + private static final int TEST_GRPC_PORT = 2183; + private static final int TEST_GRPC_SSL_PORT = 2184; private final PubSubTopicRepository pubSubTopicRepository = new PubSubTopicRepository(); + private VeniceControllerRequestHandler requestHandler; + private ControllerRequestHandlerDependencies mockDependencies; + private Admin mockAdmin; + + @BeforeMethod(alwaysRun = true) + public void setUp() { + mockAdmin = mock(VeniceParentHelixAdmin.class); + mockDependencies = mock(ControllerRequestHandlerDependencies.class); + doReturn(mockAdmin).when(mockDependencies).getAdmin(); + requestHandler = new VeniceControllerRequestHandler(mockDependencies); + } + @Test public void testGetLeaderController() throws Exception { - Admin mockAdmin = mock(VeniceParentHelixAdmin.class); doReturn(true).when(mockAdmin).isLeaderControllerFor(anyString()); - Instance leaderController = new Instance(TEST_NODE_ID, TEST_HOST, TEST_PORT, TEST_SSL_PORT); + Instance leaderController = + new Instance(TEST_NODE_ID, TEST_HOST, TEST_PORT, TEST_SSL_PORT, TEST_GRPC_PORT, TEST_GRPC_SSL_PORT); + doReturn(leaderController).when(mockAdmin).getLeaderController(anyString()); Request request = mock(Request.class); doReturn(TEST_CLUSTER).when(request).queryParams(eq(ControllerApiConstants.CLUSTER)); - Route leaderControllerRoute = - new ControllerRoutes(false, Optional.empty(), pubSubTopicRepository).getLeaderController(mockAdmin); + Route leaderControllerRoute = new ControllerRoutes(false, Optional.empty(), pubSubTopicRepository, requestHandler) + .getLeaderController(mockAdmin); LeaderControllerResponse leaderControllerResponse = OBJECT_MAPPER.readValue( leaderControllerRoute.handle(request, mock(Response.class)).toString(), LeaderControllerResponse.class); assertEquals(leaderControllerResponse.getCluster(), TEST_CLUSTER); assertEquals(leaderControllerResponse.getUrl(), "http://" + TEST_HOST + ":" + TEST_PORT); assertEquals(leaderControllerResponse.getSecureUrl(), "https://" + TEST_HOST + ":" + TEST_SSL_PORT); + assertEquals(leaderControllerResponse.getGrpcUrl(), TEST_HOST + ":" + TEST_GRPC_PORT); + assertEquals(leaderControllerResponse.getSecureGrpcUrl(), TEST_HOST + ":" + TEST_GRPC_SSL_PORT); + + when(mockDependencies.isSslEnabled()).thenReturn(true); + requestHandler = new VeniceControllerRequestHandler(mockDependencies); - Route leaderControllerSslRoute = - new ControllerRoutes(true, Optional.empty(), pubSubTopicRepository).getLeaderController(mockAdmin); + Route leaderControllerSslRoute = new ControllerRoutes(true, Optional.empty(), pubSubTopicRepository, requestHandler) + .getLeaderController(mockAdmin); LeaderControllerResponse leaderControllerResponseSsl = OBJECT_MAPPER.readValue( leaderControllerSslRoute.handle(request, mock(Response.class)).toString(), LeaderControllerResponse.class); assertEquals(leaderControllerResponseSsl.getCluster(), TEST_CLUSTER); assertEquals(leaderControllerResponseSsl.getUrl(), "https://" + TEST_HOST + ":" + TEST_SSL_PORT); assertEquals(leaderControllerResponseSsl.getSecureUrl(), "https://" + TEST_HOST + ":" + TEST_SSL_PORT); + assertEquals(leaderControllerResponse.getGrpcUrl(), TEST_HOST + ":" + TEST_GRPC_PORT); + assertEquals(leaderControllerResponse.getSecureGrpcUrl(), TEST_HOST + ":" + TEST_GRPC_SSL_PORT); // Controller doesn't support SSL - Instance leaderNonSslController = new Instance(TEST_NODE_ID, TEST_HOST, TEST_PORT, TEST_PORT); + Instance leaderNonSslController = + new Instance(TEST_NODE_ID, TEST_HOST, TEST_PORT, TEST_PORT, TEST_GRPC_PORT, TEST_GRPC_SSL_PORT); doReturn(leaderNonSslController).when(mockAdmin).getLeaderController(anyString()); LeaderControllerResponse leaderControllerNonSslResponse = OBJECT_MAPPER.readValue( @@ -78,11 +104,14 @@ public void testGetLeaderController() throws Exception { assertEquals(leaderControllerNonSslResponse.getCluster(), TEST_CLUSTER); assertEquals(leaderControllerNonSslResponse.getUrl(), "http://" + TEST_HOST + ":" + TEST_PORT); assertEquals(leaderControllerNonSslResponse.getSecureUrl(), null); + assertEquals(leaderControllerNonSslResponse.getGrpcUrl(), TEST_HOST + ":" + TEST_GRPC_PORT); + assertEquals(leaderControllerNonSslResponse.getSecureGrpcUrl(), TEST_HOST + ":" + TEST_GRPC_SSL_PORT); } @Test public void testGetAggregatedHealthStatus() throws Exception { - ControllerRoutes controllerRoutes = new ControllerRoutes(false, Optional.empty(), pubSubTopicRepository); + ControllerRoutes controllerRoutes = + new ControllerRoutes(false, Optional.empty(), pubSubTopicRepository, requestHandler); Admin mockAdmin = mock(VeniceParentHelixAdmin.class); List instanceList = Arrays.asList("instance1_5000", "instance2_5000"); diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/CreateStoreTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/CreateStoreTest.java index ebd3cd34d92..b0fdc408c57 100644 --- a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/CreateStoreTest.java +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/CreateStoreTest.java @@ -15,10 +15,13 @@ import com.linkedin.venice.HttpConstants; import com.linkedin.venice.controller.Admin; +import com.linkedin.venice.controller.ControllerRequestHandlerDependencies; +import com.linkedin.venice.controller.VeniceParentHelixAdmin; import com.linkedin.venice.utils.Utils; import java.util.HashMap; import java.util.Optional; import org.apache.commons.httpclient.HttpStatus; +import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; import spark.QueryParamsMap; import spark.Request; @@ -27,19 +30,29 @@ public class CreateStoreTest { - private static String clusterName = Utils.getUniqueString("test-cluster"); + private static final String CLUSTER_NAME = Utils.getUniqueString("test-cluster"); + + private VeniceControllerRequestHandler requestHandler; + private Admin mockAdmin; + + @BeforeMethod + public void setUp() { + mockAdmin = mock(VeniceParentHelixAdmin.class); + ControllerRequestHandlerDependencies dependencies = mock(ControllerRequestHandlerDependencies.class); + doReturn(mockAdmin).when(dependencies).getAdmin(); + requestHandler = new VeniceControllerRequestHandler(dependencies); + } @Test public void testCreateStoreWhenThrowsNPEInternally() throws Exception { - Admin admin = mock(Admin.class); Request request = mock(Request.class); Response response = mock(Response.class); String fakeMessage = "fake_message"; - doReturn(true).when(admin).isLeaderControllerFor(clusterName); + doReturn(true).when(mockAdmin).isLeaderControllerFor(CLUSTER_NAME); // Throws NPE here - doThrow(new NullPointerException(fakeMessage)).when(admin) + doThrow(new NullPointerException(fakeMessage)).when(mockAdmin) .createStore(any(), any(), any(), any(), any(), anyBoolean(), any()); QueryParamsMap paramsMap = mock(QueryParamsMap.class); @@ -47,88 +60,85 @@ public void testCreateStoreWhenThrowsNPEInternally() throws Exception { doReturn(paramsMap).when(request).queryMap(); doReturn(NEW_STORE.getPath()).when(request).pathInfo(); - doReturn(clusterName).when(request).queryParams(CLUSTER); + doReturn(CLUSTER_NAME).when(request).queryParams(CLUSTER); doReturn("test-store").when(request).queryParams(NAME); doReturn("fake-owner").when(request).queryParams(OWNER); doReturn("\"long\"").when(request).queryParams(KEY_SCHEMA); doReturn("\"string\"").when(request).queryParams(VALUE_SCHEMA); CreateStore createStoreRoute = new CreateStore(false, Optional.empty()); - Route createStoreRouter = createStoreRoute.createStore(admin); + Route createStoreRouter = createStoreRoute.createStore(mockAdmin, requestHandler); createStoreRouter.handle(request, response); verify(response).status(HttpStatus.SC_INTERNAL_SERVER_ERROR); } @Test(expectedExceptions = Error.class) public void testCreateStoreWhenThrowsError() throws Exception { - Admin admin = mock(Admin.class); Request request = mock(Request.class); Response response = mock(Response.class); String fakeMessage = "fake_message"; - doReturn(true).when(admin).isLeaderControllerFor(clusterName); + doReturn(true).when(mockAdmin).isLeaderControllerFor(CLUSTER_NAME); // Throws NPE here - doThrow(new Error(fakeMessage)).when(admin).createStore(any(), any(), any(), any(), any(), anyBoolean(), any()); + doThrow(new Error(fakeMessage)).when(mockAdmin).createStore(any(), any(), any(), any(), any(), anyBoolean(), any()); QueryParamsMap paramsMap = mock(QueryParamsMap.class); doReturn(new HashMap<>()).when(paramsMap).toMap(); doReturn(paramsMap).when(request).queryMap(); doReturn(NEW_STORE.getPath()).when(request).pathInfo(); - doReturn(clusterName).when(request).queryParams(CLUSTER); + doReturn(CLUSTER_NAME).when(request).queryParams(CLUSTER); doReturn("test-store").when(request).queryParams(NAME); doReturn("fake-owner").when(request).queryParams(OWNER); doReturn("\"long\"").when(request).queryParams(KEY_SCHEMA); doReturn("\"string\"").when(request).queryParams(VALUE_SCHEMA); CreateStore createStoreRoute = new CreateStore(false, Optional.empty()); - Route createStoreRouter = createStoreRoute.createStore(admin); + Route createStoreRouter = createStoreRoute.createStore(mockAdmin, requestHandler); createStoreRouter.handle(request, response); } @Test public void testCreateStoreWhenSomeParamNotPresent() throws Exception { - Admin admin = mock(Admin.class); Request request = mock(Request.class); Response response = mock(Response.class); - doReturn(true).when(admin).isLeaderControllerFor(clusterName); + doReturn(true).when(mockAdmin).isLeaderControllerFor(CLUSTER_NAME); QueryParamsMap paramsMap = mock(QueryParamsMap.class); doReturn(new HashMap<>()).when(paramsMap).toMap(); doReturn(paramsMap).when(request).queryMap(); doReturn(NEW_STORE.getPath()).when(request).pathInfo(); - doReturn(clusterName).when(request).queryParams(CLUSTER); + doReturn(CLUSTER_NAME).when(request).queryParams(CLUSTER); CreateStore createStoreRoute = new CreateStore(false, Optional.empty()); - Route createStoreRouter = createStoreRoute.createStore(admin); + Route createStoreRouter = createStoreRoute.createStore(mockAdmin, requestHandler); createStoreRouter.handle(request, response); verify(response).status(HttpStatus.SC_BAD_REQUEST); } @Test public void testCreateStoreWhenNotLeaderController() throws Exception { - Admin admin = mock(Admin.class); Request request = mock(Request.class); Response response = mock(Response.class); - doReturn(false).when(admin).isLeaderControllerFor(clusterName); + doReturn(false).when(mockAdmin).isLeaderControllerFor(CLUSTER_NAME); QueryParamsMap paramsMap = mock(QueryParamsMap.class); doReturn(new HashMap<>()).when(paramsMap).toMap(); doReturn(paramsMap).when(request).queryMap(); doReturn(NEW_STORE.getPath()).when(request).pathInfo(); - doReturn(clusterName).when(request).queryParams(CLUSTER); + doReturn(CLUSTER_NAME).when(request).queryParams(CLUSTER); doReturn("test-store").when(request).queryParams(NAME); doReturn("fake-owner").when(request).queryParams(OWNER); doReturn("\"long\"").when(request).queryParams(KEY_SCHEMA); doReturn("\"string\"").when(request).queryParams(VALUE_SCHEMA); CreateStore createStoreRoute = new CreateStore(false, Optional.empty()); - Route createStoreRouter = createStoreRoute.createStore(admin); + Route createStoreRouter = createStoreRoute.createStore(mockAdmin, requestHandler); createStoreRouter.handle(request, response); verify(response).status(HttpConstants.SC_MISDIRECTED_REQUEST); } diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/CreateVersionTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/CreateVersionTest.java index 39e48528682..54f7fec03bc 100644 --- a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/CreateVersionTest.java +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/CreateVersionTest.java @@ -4,55 +4,87 @@ import static com.linkedin.venice.VeniceConstants.CONTROLLER_SSL_CERTIFICATE_ATTRIBUTE_NAME; import static com.linkedin.venice.controller.server.CreateVersion.overrideSourceRegionAddressForIncrementalPushJob; import static com.linkedin.venice.controllerapi.ControllerApiConstants.CLUSTER; +import static com.linkedin.venice.controllerapi.ControllerApiConstants.COMPRESSION_DICTIONARY; +import static com.linkedin.venice.controllerapi.ControllerApiConstants.DEFER_VERSION_SWAP; import static com.linkedin.venice.controllerapi.ControllerApiConstants.HOSTNAME; +import static com.linkedin.venice.controllerapi.ControllerApiConstants.IS_WRITE_COMPUTE_ENABLED; import static com.linkedin.venice.controllerapi.ControllerApiConstants.NAME; +import static com.linkedin.venice.controllerapi.ControllerApiConstants.PARTITIONERS; +import static com.linkedin.venice.controllerapi.ControllerApiConstants.PUSH_IN_SORTED_ORDER; import static com.linkedin.venice.controllerapi.ControllerApiConstants.PUSH_JOB_ID; import static com.linkedin.venice.controllerapi.ControllerApiConstants.PUSH_TYPE; import static com.linkedin.venice.controllerapi.ControllerApiConstants.REPUSH_SOURCE_VERSION; +import static com.linkedin.venice.controllerapi.ControllerApiConstants.REWIND_TIME_IN_SECONDS_OVERRIDE; +import static com.linkedin.venice.controllerapi.ControllerApiConstants.SEND_START_OF_PUSH; import static com.linkedin.venice.controllerapi.ControllerApiConstants.SEPARATE_REAL_TIME_TOPIC_ENABLED; +import static com.linkedin.venice.controllerapi.ControllerApiConstants.SOURCE_GRID_FABRIC; import static com.linkedin.venice.controllerapi.ControllerApiConstants.STORE_SIZE; +import static com.linkedin.venice.controllerapi.ControllerApiConstants.TARGETED_REGIONS; import static com.linkedin.venice.controllerapi.ControllerRoute.REQUEST_TOPIC; import static com.linkedin.venice.meta.BufferReplayPolicy.REWIND_FROM_EOP; import static com.linkedin.venice.meta.DataReplicationPolicy.ACTIVE_ACTIVE; import static com.linkedin.venice.meta.DataReplicationPolicy.AGGREGATE; import static com.linkedin.venice.meta.DataReplicationPolicy.NONE; import static com.linkedin.venice.meta.DataReplicationPolicy.NON_AGGREGATE; +import static com.linkedin.venice.meta.Version.PushType.BATCH; +import static com.linkedin.venice.meta.Version.PushType.INCREMENTAL; +import static com.linkedin.venice.meta.Version.PushType.STREAM; +import static com.linkedin.venice.meta.Version.PushType.STREAM_REPROCESSING; +import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyBoolean; -import static org.mockito.Mockito.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; import static org.mockito.Mockito.doCallRealMethod; import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; import static org.testng.Assert.assertTrue; import static org.testng.Assert.expectThrows; +import static org.testng.Assert.fail; import com.fasterxml.jackson.databind.ObjectMapper; import com.linkedin.venice.acl.DynamicAccessController; +import com.linkedin.venice.compression.CompressionStrategy; import com.linkedin.venice.controller.Admin; +import com.linkedin.venice.controllerapi.RequestTopicForPushRequest; import com.linkedin.venice.controllerapi.VersionCreationResponse; import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.exceptions.VeniceHttpException; +import com.linkedin.venice.exceptions.VeniceUnsupportedOperationException; +import com.linkedin.venice.meta.DataReplicationPolicy; +import com.linkedin.venice.meta.HybridStoreConfig; import com.linkedin.venice.meta.HybridStoreConfigImpl; import com.linkedin.venice.meta.OfflinePushStrategy; +import com.linkedin.venice.meta.PartitionerConfig; import com.linkedin.venice.meta.PersistenceType; import com.linkedin.venice.meta.ReadStrategy; import com.linkedin.venice.meta.RoutingStrategy; import com.linkedin.venice.meta.Store; import com.linkedin.venice.meta.Version; -import com.linkedin.venice.meta.Version.PushType; import com.linkedin.venice.meta.VersionImpl; import com.linkedin.venice.meta.ZKStore; import com.linkedin.venice.utils.DataProviderUtils; import com.linkedin.venice.utils.ObjectMapperFactory; import com.linkedin.venice.utils.Utils; +import com.linkedin.venice.utils.lazy.Lazy; import java.security.cert.X509Certificate; +import java.util.Arrays; +import java.util.Collections; import java.util.HashMap; +import java.util.HashSet; import java.util.Map; import java.util.Optional; +import java.util.Set; import javax.security.auth.x500.X500Principal; import javax.servlet.http.HttpServletRequest; +import org.apache.http.HttpStatus; import org.testng.Assert; import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; @@ -95,7 +127,7 @@ public void setUp() { queryMap.put(NAME, new String[] { STORE_NAME }); queryMap.put(STORE_SIZE, new String[] { "0" }); queryMap.put(REPUSH_SOURCE_VERSION, new String[] { "0" }); - queryMap.put(PUSH_TYPE, new String[] { PushType.INCREMENTAL.name() }); + queryMap.put(PUSH_TYPE, new String[] { INCREMENTAL.name() }); queryMap.put(PUSH_JOB_ID, new String[] { JOB_ID }); queryMap.put(HOSTNAME, new String[] { "localhost" }); @@ -168,7 +200,7 @@ public void testRequestTopicForHybridIncPushEnabled( JOB_ID, 0, 0, - PushType.INCREMENTAL, + INCREMENTAL, false, false, null, @@ -227,7 +259,7 @@ public void testRequestTopicForIncPushReturnsErrorWhenStoreIsNotHybridAndIncPush JOB_ID, 0, 0, - PushType.INCREMENTAL, + INCREMENTAL, false, false, null, @@ -253,7 +285,7 @@ public void testRequestTopicForIncPushReturnsErrorWhenStoreIsNotHybridAndIncPush OBJECT_MAPPER.readValue(result.toString(), VersionCreationResponse.class); assertTrue(versionCreateResponse.isError()); assertTrue(versionCreateResponse.getError().contains("which does not have hybrid mode enabled")); - Assert.assertNull(versionCreateResponse.getKafkaTopic()); + assertNull(versionCreateResponse.getKafkaTopic()); } @Test @@ -281,7 +313,7 @@ public void testRequestTopicForIncPushCanUseEmergencyRegionWhenItIsSet() throws JOB_ID, 0, 0, - PushType.INCREMENTAL, + INCREMENTAL, false, false, null, @@ -336,39 +368,87 @@ public void testOverrideSourceRegionAddressForIncrementalPushJob() { creationResponse = new VersionCreationResponse(); creationResponse.setKafkaBootstrapServers("default.src.region.com"); doReturn(Optional.empty()).when(admin).getAggregateRealTimeTopicSource(CLUSTER_NAME); - overrideSourceRegionAddressForIncrementalPushJob(admin, creationResponse, CLUSTER_NAME, null, null, false, true); + overrideSourceRegionAddressForIncrementalPushJob( + admin, + creationResponse, + CLUSTER_NAME, + STORE_NAME, + null, + null, + false, + true); assertEquals(creationResponse.getKafkaBootstrapServers(), "default.src.region.com"); // AA-all-region is disabled & NR is enabled * AGG RT address is set creationResponse = new VersionCreationResponse(); creationResponse.setKafkaBootstrapServers("default.src.region.com"); doReturn(Optional.of("agg.rt.region.com")).when(admin).getAggregateRealTimeTopicSource(CLUSTER_NAME); - overrideSourceRegionAddressForIncrementalPushJob(admin, creationResponse, CLUSTER_NAME, null, null, false, true); + overrideSourceRegionAddressForIncrementalPushJob( + admin, + creationResponse, + CLUSTER_NAME, + STORE_NAME, + null, + null, + false, + true); assertEquals(creationResponse.getKafkaBootstrapServers(), "agg.rt.region.com"); // AA-all-region and NR are disabled creationResponse = new VersionCreationResponse(); creationResponse.setKafkaBootstrapServers("default.src.region.com"); - overrideSourceRegionAddressForIncrementalPushJob(admin, creationResponse, CLUSTER_NAME, null, null, false, false); + overrideSourceRegionAddressForIncrementalPushJob( + admin, + creationResponse, + CLUSTER_NAME, + STORE_NAME, + null, + null, + false, + false); assertEquals(creationResponse.getKafkaBootstrapServers(), "default.src.region.com"); // AA-all-region is enabled and NR is disabled creationResponse = new VersionCreationResponse(); creationResponse.setKafkaBootstrapServers("default.src.region.com"); - overrideSourceRegionAddressForIncrementalPushJob(admin, creationResponse, CLUSTER_NAME, null, null, true, false); + overrideSourceRegionAddressForIncrementalPushJob( + admin, + creationResponse, + CLUSTER_NAME, + STORE_NAME, + null, + null, + true, + false); assertEquals(creationResponse.getKafkaBootstrapServers(), "default.src.region.com"); // AA-all-region and NR are enabled AND emergencySourceRegion and pushJobSourceGridFabric are null creationResponse = new VersionCreationResponse(); creationResponse.setKafkaBootstrapServers("default.src.region.com"); - overrideSourceRegionAddressForIncrementalPushJob(admin, creationResponse, CLUSTER_NAME, null, null, true, true); + overrideSourceRegionAddressForIncrementalPushJob( + admin, + creationResponse, + CLUSTER_NAME, + STORE_NAME, + null, + null, + true, + true); assertEquals(creationResponse.getKafkaBootstrapServers(), "default.src.region.com"); // AA-all-region and NR are enabled AND emergencySourceRegion is not set but pushJobSourceGridFabric is provided creationResponse = new VersionCreationResponse(); creationResponse.setKafkaBootstrapServers("default.src.region.com"); doReturn("vpj.src.region.com").when(admin).getNativeReplicationKafkaBootstrapServerAddress("dc-vpj"); - overrideSourceRegionAddressForIncrementalPushJob(admin, creationResponse, CLUSTER_NAME, null, "dc-vpj", true, true); + overrideSourceRegionAddressForIncrementalPushJob( + admin, + creationResponse, + CLUSTER_NAME, + STORE_NAME, + null, + "dc-vpj", + true, + true); assertEquals(creationResponse.getKafkaBootstrapServers(), "vpj.src.region.com"); // AA-all-region and NR are enabled AND emergencySourceRegion is set and pushJobSourceGridFabric is provided @@ -379,6 +459,7 @@ public void testOverrideSourceRegionAddressForIncrementalPushJob() { admin, creationResponse, CLUSTER_NAME, + STORE_NAME, "dc-e", "dc-vpj", true, @@ -389,7 +470,15 @@ public void testOverrideSourceRegionAddressForIncrementalPushJob() { creationResponse = new VersionCreationResponse(); creationResponse.setKafkaBootstrapServers("emergency.src.region.com"); doReturn("emergency.src.region.com").when(admin).getNativeReplicationKafkaBootstrapServerAddress("dc-e"); - overrideSourceRegionAddressForIncrementalPushJob(admin, creationResponse, CLUSTER_NAME, "dc-e", null, true, true); + overrideSourceRegionAddressForIncrementalPushJob( + admin, + creationResponse, + CLUSTER_NAME, + STORE_NAME, + "dc-e", + null, + true, + true); assertEquals(creationResponse.getKafkaBootstrapServers(), "emergency.src.region.com"); } @@ -398,7 +487,15 @@ public void testOverrideSourceRegionAddressForIncrementalPushJobWhenOverrideRegi VersionCreationResponse creationResponse = new VersionCreationResponse(); creationResponse.setKafkaBootstrapServers("default.src.region.com"); doReturn(null).when(admin).getNativeReplicationKafkaBootstrapServerAddress("dc1"); - overrideSourceRegionAddressForIncrementalPushJob(admin, creationResponse, CLUSTER_NAME, "dc1", null, true, true); + overrideSourceRegionAddressForIncrementalPushJob( + admin, + creationResponse, + CLUSTER_NAME, + STORE_NAME, + "dc1", + null, + true, + true); } @Test @@ -408,35 +505,35 @@ public void testValidatePushTypeForStreamPushType() { // push type is STREAM and store is not hybrid Store store1 = mock(Store.class); when(store1.isHybrid()).thenReturn(false); - Exception e = expectThrows(VeniceException.class, () -> createVersion.validatePushType(PushType.STREAM, store1)); + Exception e = expectThrows(VeniceException.class, () -> createVersion.validatePushType(STREAM, store1)); assertTrue(e.getMessage().contains("which is not configured to be a hybrid store")); // push type is STREAM and store is AA enabled hybrid Store store2 = mock(Store.class); when(store2.isHybrid()).thenReturn(true); when(store2.isActiveActiveReplicationEnabled()).thenReturn(true); - createVersion.validatePushType(PushType.STREAM, store2); + createVersion.validatePushType(STREAM, store2); // push type is STREAM and store is not AA enabled hybrid but has NON_AGGREGATE replication policy Store store3 = mock(Store.class); when(store3.isHybrid()).thenReturn(true); when(store3.isActiveActiveReplicationEnabled()).thenReturn(false); when(store3.getHybridStoreConfig()).thenReturn(new HybridStoreConfigImpl(0, 1, 0, NON_AGGREGATE, REWIND_FROM_EOP)); - createVersion.validatePushType(PushType.STREAM, store3); + createVersion.validatePushType(STREAM, store3); // push type is STREAM and store is not AA enabled hybrid but has AGGREGATE replication policy Store store4 = mock(Store.class); when(store4.isHybrid()).thenReturn(true); when(store4.isActiveActiveReplicationEnabled()).thenReturn(false); when(store4.getHybridStoreConfig()).thenReturn(new HybridStoreConfigImpl(0, 1, 0, AGGREGATE, REWIND_FROM_EOP)); - createVersion.validatePushType(PushType.STREAM, store4); + createVersion.validatePushType(STREAM, store4); // push type is STREAM and store is not AA enabled hybrid but has NONE replication policy Store store5 = mock(Store.class); when(store5.isHybrid()).thenReturn(true); when(store5.isActiveActiveReplicationEnabled()).thenReturn(false); when(store5.getHybridStoreConfig()).thenReturn(new HybridStoreConfigImpl(0, 1, 0, NONE, REWIND_FROM_EOP)); - Exception e5 = expectThrows(VeniceException.class, () -> createVersion.validatePushType(PushType.STREAM, store5)); + Exception e5 = expectThrows(VeniceException.class, () -> createVersion.validatePushType(STREAM, store5)); assertTrue(e5.getMessage().contains("which is configured to have a hybrid data replication policy")); // push type is STREAM and store is not AA enabled hybrid but has ACTIVE_ACTIVE replication policy @@ -444,7 +541,7 @@ public void testValidatePushTypeForStreamPushType() { when(store6.isHybrid()).thenReturn(true); when(store6.isActiveActiveReplicationEnabled()).thenReturn(false); when(store6.getHybridStoreConfig()).thenReturn(new HybridStoreConfigImpl(0, 1, 0, ACTIVE_ACTIVE, REWIND_FROM_EOP)); - Exception e6 = expectThrows(VeniceException.class, () -> createVersion.validatePushType(PushType.STREAM, store6)); + Exception e6 = expectThrows(VeniceException.class, () -> createVersion.validatePushType(STREAM, store6)); assertTrue(e6.getMessage().contains("which is configured to have a hybrid data replication policy")); } @@ -455,16 +552,547 @@ public void testValidatePushTypeForIncrementalPushPushType() { // push type is INCREMENTAL and store is not hybrid Store store1 = mock(Store.class); when(store1.isHybrid()).thenReturn(false); - Exception e = - expectThrows(VeniceException.class, () -> createVersion.validatePushType(PushType.INCREMENTAL, store1)); + Exception e = expectThrows(VeniceException.class, () -> createVersion.validatePushType(INCREMENTAL, store1)); assertTrue(e.getMessage().contains("which does not have hybrid mode enabled")); // push type is INCREMENTAL and store is hybrid but incremental push is not enabled Store store2 = mock(Store.class); when(store2.isHybrid()).thenReturn(true); when(store2.isIncrementalPushEnabled()).thenReturn(false); - Exception e2 = - expectThrows(VeniceException.class, () -> createVersion.validatePushType(PushType.INCREMENTAL, store2)); + Exception e2 = expectThrows(VeniceException.class, () -> createVersion.validatePushType(INCREMENTAL, store2)); assertTrue(e2.getMessage().contains("which does not have incremental push enabled")); } + + @Test + public void testExtractOptionalParamsFromRequestTopicForPushingRequest() { + // Test case 1: Default values + Request mockRequest = mock(Request.class); + doCallRealMethod().when(mockRequest).queryParamOrDefault(anyString(), anyString()); + doReturn(null).when(mockRequest).queryParams(any()); + + RequestTopicForPushRequest requestDetails = new RequestTopicForPushRequest(CLUSTER_NAME, STORE_NAME, BATCH, JOB_ID); + + CreateVersion.extractOptionalParamsFromRequestTopicRequest(mockRequest, requestDetails, false); + + assertNotNull(requestDetails.getPartitioners(), "Default partitioners should not be null"); + assertTrue(requestDetails.getPartitioners().isEmpty(), "Default partitioners should be empty"); + assertFalse(requestDetails.isSendStartOfPush(), "Default sendStartOfPush should be false"); + assertFalse(requestDetails.isSorted(), "Default sorted should be false"); + assertFalse(requestDetails.isWriteComputeEnabled(), "Default writeComputeEnabled should be false"); + assertEquals( + requestDetails.getRewindTimeInSecondsOverride(), + -1L, + "Default rewindTimeInSecondsOverride should be -1"); + assertFalse(requestDetails.isDeferVersionSwap(), "Default deferVersionSwap should be false"); + assertNull(requestDetails.getTargetedRegions(), "Default targetedRegions should be null"); + assertEquals(requestDetails.getRepushSourceVersion(), -1, "Default repushSourceVersion should be -1"); + assertNull(requestDetails.getSourceGridFabric(), "Default sourceGridFabric should be null"); + assertNull(requestDetails.getCompressionDictionary(), "Default compressionDictionary should be null"); + assertNull(requestDetails.getCertificateInRequest(), "Default certificateInRequest should be null"); + + // Test case 2: All optional parameters are set + mockRequest = mock(Request.class); + doCallRealMethod().when(mockRequest).queryParamOrDefault(any(), any()); + String customPartitioners = "f.q.c.n.P1,f.q.c.n.P2"; + Set expectedPartitioners = new HashSet<>(Arrays.asList("f.q.c.n.P1", "f.q.c.n.P2")); + + when(mockRequest.queryParams(eq(PARTITIONERS))).thenReturn(customPartitioners); + when(mockRequest.queryParams(SEND_START_OF_PUSH)).thenReturn("true"); + when(mockRequest.queryParams(PUSH_IN_SORTED_ORDER)).thenReturn("true"); + when(mockRequest.queryParams(IS_WRITE_COMPUTE_ENABLED)).thenReturn("true"); + when(mockRequest.queryParams(REWIND_TIME_IN_SECONDS_OVERRIDE)).thenReturn("120"); + when(mockRequest.queryParams(DEFER_VERSION_SWAP)).thenReturn("true"); + when(mockRequest.queryParams(TARGETED_REGIONS)).thenReturn("region-1"); + when(mockRequest.queryParams(REPUSH_SOURCE_VERSION)).thenReturn("5"); + when(mockRequest.queryParams(SOURCE_GRID_FABRIC)).thenReturn("grid-fabric"); + when(mockRequest.queryParams(COMPRESSION_DICTIONARY)).thenReturn("XYZ"); + + requestDetails = new RequestTopicForPushRequest(CLUSTER_NAME, STORE_NAME, BATCH, JOB_ID); + + CreateVersion.extractOptionalParamsFromRequestTopicRequest(mockRequest, requestDetails, false); + + assertEquals(requestDetails.getPartitioners(), expectedPartitioners); + assertTrue(requestDetails.isSendStartOfPush()); + assertTrue(requestDetails.isSorted()); + assertTrue(requestDetails.isWriteComputeEnabled()); + assertEquals(requestDetails.getRewindTimeInSecondsOverride(), 120L); + assertTrue(requestDetails.isDeferVersionSwap()); + assertEquals(requestDetails.getTargetedRegions(), "region-1"); + assertEquals(requestDetails.getRepushSourceVersion(), 5); + assertEquals(requestDetails.getSourceGridFabric(), "grid-fabric"); + assertEquals(requestDetails.getCompressionDictionary(), "XYZ"); + + // Test case 3: check that the certificate is set in the request details when access control is enabled + HttpServletRequest mockHttpServletRequest = mock(HttpServletRequest.class); + X509Certificate[] mockCertificates = { mock(X509Certificate.class) }; + when(mockHttpServletRequest.getAttribute(CONTROLLER_SSL_CERTIFICATE_ATTRIBUTE_NAME)).thenReturn(mockCertificates); + when(mockRequest.raw()).thenReturn(mockHttpServletRequest); + CreateVersion.extractOptionalParamsFromRequestTopicRequest(mockRequest, requestDetails, true); + assertEquals(requestDetails.getCertificateInRequest(), mockCertificates[0]); + + // Test case 4: Invalid values for optional parameters + when(mockRequest.queryParams(SEND_START_OF_PUSH)).thenReturn("notBoolean"); + when(mockRequest.queryParams(REWIND_TIME_IN_SECONDS_OVERRIDE)).thenReturn("invalidLong"); + + requestDetails = new RequestTopicForPushRequest(CLUSTER_NAME, STORE_NAME, BATCH, JOB_ID); + Request finalMockRequest = mockRequest; + RequestTopicForPushRequest finalRequestDetails = requestDetails; + VeniceHttpException e = expectThrows( + VeniceHttpException.class, + () -> CreateVersion.extractOptionalParamsFromRequestTopicRequest(finalMockRequest, finalRequestDetails, false)); + assertEquals(e.getHttpStatusCode(), HttpStatus.SC_BAD_REQUEST); + } + + @Test + public void testVerifyAndConfigurePartitionerSettings() { + CreateVersion createVersion = new CreateVersion(true, Optional.of(accessClient), false, false); + + VersionCreationResponse response = new VersionCreationResponse(); + PartitionerConfig storePartitionerConfig = mock(PartitionerConfig.class); + when(storePartitionerConfig.getPartitionerClass()).thenReturn("f.q.c.n.DefaultPartitioner"); + + // Test Case 1: Null partitionersFromRequest (should pass) + try { + createVersion.verifyAndConfigurePartitionerSettings(storePartitionerConfig, null, response); + } catch (Exception e) { + fail("Null partitionersFromRequest should not throw an exception."); + } + assertEquals(response.getPartitionerClass(), "f.q.c.n.DefaultPartitioner"); + + // Test Case 2: Empty partitionersFromRequest (should pass) + response = new VersionCreationResponse(); + Set partitionersFromRequest = Collections.emptySet(); + try { + createVersion.verifyAndConfigurePartitionerSettings(storePartitionerConfig, partitionersFromRequest, response); + } catch (Exception e) { + fail("Empty partitionersFromRequest should not throw an exception."); + } + assertEquals(response.getPartitionerClass(), "f.q.c.n.DefaultPartitioner"); + + // Test Case 3: Matching partitioner in partitionersFromRequest (should pass) + response = new VersionCreationResponse(); + partitionersFromRequest = new HashSet<>(Arrays.asList("f.q.c.n.DefaultPartitioner", "f.q.c.n.CustomPartitioner")); + try { + createVersion.verifyAndConfigurePartitionerSettings(storePartitionerConfig, partitionersFromRequest, response); + } catch (Exception e) { + fail("Matching partitioner should not throw an exception."); + } + assertEquals(response.getPartitionerClass(), "f.q.c.n.DefaultPartitioner"); + + // Test Case 4: Non-matching partitioner in partitionersFromRequest (should throw exception) + final VersionCreationResponse finalResponse = new VersionCreationResponse(); + partitionersFromRequest = new HashSet<>(Collections.singletonList("f.q.c.n.CustomPartitioner")); + Set finalPartitionersFromRequest = partitionersFromRequest; + Exception e = expectThrows( + VeniceException.class, + () -> createVersion.verifyAndConfigurePartitionerSettings( + storePartitionerConfig, + finalPartitionersFromRequest, + finalResponse)); + assertTrue(e.getMessage().contains("cannot be found")); + } + + @Test + public void testDetermineResponseTopic() { + CreateVersion createVersion = new CreateVersion(true, Optional.of(accessClient), false, false); + + String storeName = "test_store"; + String vtName = Version.composeKafkaTopic(storeName, 1); + String rtName = Version.composeRealTimeTopic(storeName); + String srTopicName = Version.composeStreamReprocessingTopic(storeName, 1); + String separateRtName = Version.composeSeparateRealTimeTopic(storeName); + + RequestTopicForPushRequest request = new RequestTopicForPushRequest("v0", storeName, INCREMENTAL, "JOB_ID"); + + // Test Case: PushType.INCREMENTAL with separate real-time topic enabled + Version mockVersion1 = mock(Version.class); + when(mockVersion1.kafkaTopicName()).thenReturn(vtName); + when(mockVersion1.isSeparateRealTimeTopicEnabled()).thenReturn(true); + request.setSeparateRealTimeTopicEnabled(true); + String result1 = createVersion.determineResponseTopic(storeName, mockVersion1, request); + assertEquals(result1, separateRtName); + + // Test Case: PushType.INCREMENTAL with separate real-time topic enabled, but the request does not have the separate + // real-time topic flag + mockVersion1 = mock(Version.class); + when(mockVersion1.getStoreName()).thenReturn(storeName); + when(mockVersion1.kafkaTopicName()).thenReturn(vtName); + when(mockVersion1.isSeparateRealTimeTopicEnabled()).thenReturn(true); + request.setSeparateRealTimeTopicEnabled(false); + result1 = createVersion.determineResponseTopic(storeName, mockVersion1, request); + assertEquals(result1, rtName); + + // Test Case: PushType.INCREMENTAL without separate real-time topic enabled + Version mockVersion2 = mock(Version.class); + when(mockVersion2.getStoreName()).thenReturn(storeName); + when(mockVersion2.kafkaTopicName()).thenReturn(vtName); + when(mockVersion2.isSeparateRealTimeTopicEnabled()).thenReturn(true); + request = new RequestTopicForPushRequest("v0", storeName, INCREMENTAL, "JOB_ID"); + String result2 = createVersion.determineResponseTopic(storeName, mockVersion2, request); + assertEquals(result2, rtName); + + // Test Case: PushType.STREAM + Version mockVersion3 = mock(Version.class); + when(mockVersion3.getStoreName()).thenReturn(storeName); + when(mockVersion3.kafkaTopicName()).thenReturn(vtName); + request = new RequestTopicForPushRequest("v0", storeName, STREAM, "JOB_ID"); + String result3 = createVersion.determineResponseTopic(storeName, mockVersion3, request); + assertEquals(result3, rtName); + + // Test Case: PushType.STREAM_REPROCESSING + Version mockVersion4 = mock(Version.class); + when(mockVersion4.getStoreName()).thenReturn(storeName); + when(mockVersion4.kafkaTopicName()).thenReturn(vtName); + when(mockVersion4.getNumber()).thenReturn(1); + request = new RequestTopicForPushRequest("v0", storeName, STREAM_REPROCESSING, "JOB_ID"); + String result4 = createVersion.determineResponseTopic(storeName, mockVersion4, request); + assertEquals(result4, srTopicName); + + // Test Case: Default case with a Kafka topic + Version mockVersion5 = mock(Version.class); + when(mockVersion5.getStoreName()).thenReturn(storeName); + when(mockVersion5.kafkaTopicName()).thenReturn(vtName); + request = new RequestTopicForPushRequest("v0", storeName, BATCH, "JOB_ID"); + String result5 = createVersion.determineResponseTopic(storeName, mockVersion5, request); + assertEquals(result5, vtName); + } + + @Test + public void testGetCompressionStrategy() { + CreateVersion createVersion = new CreateVersion(true, Optional.of(accessClient), false, false); + + // Test Case 1: Real-time topic returns NO_OP + Version mockVersion1 = mock(Version.class); + String responseTopic1 = Version.composeRealTimeTopic("test_store"); + CompressionStrategy result1 = createVersion.getCompressionStrategy(mockVersion1, responseTopic1); + assertEquals(result1, CompressionStrategy.NO_OP); + + // Test Case 2: Non-real-time topic returns version's compression strategy + Version mockVersion2 = mock(Version.class); + String responseTopic2 = Version.composeKafkaTopic("test_store", 1); + when(mockVersion2.getCompressionStrategy()).thenReturn(CompressionStrategy.GZIP); + CompressionStrategy result2 = createVersion.getCompressionStrategy(mockVersion2, responseTopic2); + assertEquals(result2, CompressionStrategy.GZIP); + } + + @Test + public void testConfigureSourceFabric() { + CreateVersion createVersion = new CreateVersion(true, Optional.of(accessClient), false, false); + + // Test Case 1: Native replication enabled and non-incremental push type + Admin mockAdmin1 = mock(Admin.class); + Version mockVersion1 = mock(Version.class); + Lazy mockLazy1 = mock(Lazy.class); + RequestTopicForPushRequest mockRequest1 = mock(RequestTopicForPushRequest.class); + VersionCreationResponse mockResponse1 = new VersionCreationResponse(); + + when(mockVersion1.isNativeReplicationEnabled()).thenReturn(true); + when(mockVersion1.getPushStreamSourceAddress()).thenReturn("bootstrapServer1"); + when(mockVersion1.getNativeReplicationSourceFabric()).thenReturn("sourceFabric1"); + when(mockRequest1.getPushType()).thenReturn(BATCH); + + createVersion.configureSourceFabric(mockAdmin1, mockVersion1, mockLazy1, mockRequest1, mockResponse1); + + assertEquals(mockResponse1.getKafkaBootstrapServers(), "bootstrapServer1"); + assertEquals(mockResponse1.getKafkaSourceRegion(), "sourceFabric1"); + + // Test Case 2: Native replication enabled with null PushStreamSourceAddress + Admin mockAdmin2 = mock(Admin.class); + Version mockVersion2 = mock(Version.class); + Lazy mockLazy2 = mock(Lazy.class); + RequestTopicForPushRequest mockRequest2 = mock(RequestTopicForPushRequest.class); + VersionCreationResponse mockResponse2 = new VersionCreationResponse(); + + when(mockVersion2.isNativeReplicationEnabled()).thenReturn(true); + when(mockVersion2.getPushStreamSourceAddress()).thenReturn(null); + when(mockVersion2.getNativeReplicationSourceFabric()).thenReturn("sourceFabric2"); + when(mockRequest2.getPushType()).thenReturn(BATCH); + + createVersion.configureSourceFabric(mockAdmin2, mockVersion2, mockLazy2, mockRequest2, mockResponse2); + + assertNull(mockResponse2.getKafkaBootstrapServers()); + assertEquals(mockResponse2.getKafkaSourceRegion(), "sourceFabric2"); + + // Test Case 3: Incremental push with parent admin and override source region + Admin mockAdmin3 = mock(Admin.class); + Version mockVersion3 = mock(Version.class); + Lazy mockLazy3 = mock(Lazy.class); + RequestTopicForPushRequest mockRequest3 = mock(RequestTopicForPushRequest.class); + VersionCreationResponse mockResponse3 = new VersionCreationResponse(); + + when(mockAdmin3.isParent()).thenReturn(true); + when(mockVersion3.isNativeReplicationEnabled()).thenReturn(true); + when(mockRequest3.getPushType()).thenReturn(INCREMENTAL); + when(mockRequest3.getClusterName()).thenReturn("testCluster"); + when(mockRequest3.getStoreName()).thenReturn("testStore"); + when(mockRequest3.getEmergencySourceRegion()).thenReturn("emergencyRegion"); + when(mockRequest3.getSourceGridFabric()).thenReturn("gridFabric"); + when(mockLazy3.get()).thenReturn(true); + + when(mockAdmin3.getNativeReplicationKafkaBootstrapServerAddress("emergencyRegion")) + .thenReturn("emergencyRegionAddress"); + + createVersion.configureSourceFabric(mockAdmin3, mockVersion3, mockLazy3, mockRequest3, mockResponse3); + + assertEquals(mockResponse3.getKafkaBootstrapServers(), "emergencyRegionAddress"); + + // No specific assertions here since `overrideSourceRegionAddressForIncrementalPushJob` is mocked, + // but we can verify if the mock was called with appropriate parameters. + verify(mockAdmin3, times(1)).isParent(); + } + + @Test + public void testHandleStreamPushTypeInParentController() { + Admin admin = mock(Admin.class); + Store store = mock(Store.class); + when(store.getName()).thenReturn(STORE_NAME); + HybridStoreConfig hybridStoreConfig = mock(HybridStoreConfig.class); + when(store.getHybridStoreConfig()).thenReturn(hybridStoreConfig); + RequestTopicForPushRequest request = new RequestTopicForPushRequest("CLUSTER_NAME", STORE_NAME, STREAM, "JOB_ID"); + VersionCreationResponse response = new VersionCreationResponse(); + + // Case 1: Parent region; With stream pushes disabled + when(admin.isParent()).thenReturn(true); + CreateVersion createVersionNotOk = new CreateVersion(true, Optional.of(accessClient), false, true); + VeniceException ex1 = expectThrows( + VeniceException.class, + () -> createVersionNotOk.handleStreamPushType(admin, store, request, response, Lazy.of(() -> false))); + assertTrue( + ex1.getMessage().contains("Write operations to the parent region are not permitted with push type: STREAM")); + + CreateVersion createVersionOk = new CreateVersion(true, Optional.of(accessClient), false, false); + + // Case 2: Parent region; Non-aggregate mode in parent with no AA replication + when(admin.isParent()).thenReturn(true); + when(store.getHybridStoreConfig().getDataReplicationPolicy()).thenReturn(NON_AGGREGATE); + VeniceException ex2 = expectThrows( + VeniceException.class, + () -> createVersionOk.handleStreamPushType(admin, store, request, response, Lazy.of(() -> false))); + assertTrue(ex2.getMessage().contains("Store is not in aggregate mode!")); + + // Case 3: Parent region; Non-aggregate mode but AA replication enabled and no hybrid version + when(admin.isParent()).thenReturn(true); + when(store.getHybridStoreConfig().getDataReplicationPolicy()).thenReturn(NON_AGGREGATE); + when(store.isActiveActiveReplicationEnabled()).thenReturn(true); + when(admin.getReferenceVersionForStreamingWrites(anyString(), anyString(), anyString())).thenReturn(null); + VeniceException ex3 = expectThrows( + VeniceException.class, + () -> createVersionOk.handleStreamPushType(admin, store, request, response, Lazy.of(() -> true))); + assertTrue(ex3.getMessage().contains("No hybrid version found for store"), "Got: " + ex3.getMessage()); + + // Case 4: Parent region; Aggregate mode but no hybrid version + when(admin.isParent()).thenReturn(true); + when(store.getHybridStoreConfig().getDataReplicationPolicy()).thenReturn(AGGREGATE); + when(store.isActiveActiveReplicationEnabled()).thenReturn(false); + when(admin.getReferenceVersionForStreamingWrites(anyString(), anyString(), anyString())).thenReturn(null); + VeniceException ex4 = expectThrows( + VeniceException.class, + () -> createVersionOk.handleStreamPushType(admin, store, request, response, Lazy.of(() -> true))); + assertTrue(ex4.getMessage().contains("No hybrid version found for store"), "Got: " + ex4.getMessage()); + + // Case 5: Parent region; Aggregate mode and there is a hybrid version + Version mockVersion = mock(Version.class); + when(mockVersion.getPartitionCount()).thenReturn(42); + when(admin.isParent()).thenReturn(true); + when(store.getHybridStoreConfig().getDataReplicationPolicy()).thenReturn(AGGREGATE); + when(admin.getReferenceVersionForStreamingWrites(anyString(), anyString(), anyString())).thenReturn(mockVersion); + createVersionOk.handleStreamPushType(admin, store, request, response, Lazy.of(() -> true)); + assertEquals(response.getPartitions(), 42); + assertEquals(response.getCompressionStrategy(), CompressionStrategy.NO_OP); + assertEquals(response.getKafkaTopic(), Version.composeRealTimeTopic(STORE_NAME)); + } + + @Test + public void testHandleStreamPushTypeInChildController() { + Admin admin = mock(Admin.class); + Store store = mock(Store.class); + when(store.getName()).thenReturn(STORE_NAME); + HybridStoreConfig hybridStoreConfig = mock(HybridStoreConfig.class); + when(store.getHybridStoreConfig()).thenReturn(hybridStoreConfig); + RequestTopicForPushRequest request = new RequestTopicForPushRequest("CLUSTER_NAME", STORE_NAME, STREAM, "JOB_ID"); + VersionCreationResponse response = new VersionCreationResponse(); + CreateVersion createVersionOk = new CreateVersion(true, Optional.of(accessClient), false, false); + + // Case 1: Child region; Aggregate mode in child and AA not enabled + when(admin.isParent()).thenReturn(false); + when(store.getHybridStoreConfig().getDataReplicationPolicy()).thenReturn(DataReplicationPolicy.AGGREGATE); + when(store.isActiveActiveReplicationEnabled()).thenReturn(false); + VeniceException ex5 = expectThrows( + VeniceException.class, + () -> createVersionOk.handleStreamPushType(admin, store, request, response, Lazy.of(() -> false))); + assertTrue(ex5.getMessage().contains("Store is in aggregate mode and AA is not enabled")); + + // Case 2: Child region; Aggregate mode but AA is enabled in all regions but no hybrid version + when(admin.isParent()).thenReturn(false); + when(store.isActiveActiveReplicationEnabled()).thenReturn(true); + when(store.getHybridStoreConfig().getDataReplicationPolicy()).thenReturn(DataReplicationPolicy.AGGREGATE); + when(admin.getReferenceVersionForStreamingWrites(anyString(), anyString(), anyString())).thenReturn(null); + VeniceException ex6 = expectThrows( + VeniceException.class, + () -> createVersionOk.handleStreamPushType(admin, store, request, response, Lazy.of(() -> true))); + assertTrue(ex6.getMessage().contains("No hybrid version found"), "Got: " + ex6.getMessage()); + + // Case 3: Child region; Non-aggregate mode but no hybrid version + when(admin.isParent()).thenReturn(false); + when(store.getHybridStoreConfig().getDataReplicationPolicy()).thenReturn(DataReplicationPolicy.NON_AGGREGATE); + when(admin.getReferenceVersionForStreamingWrites(anyString(), anyString(), anyString())).thenReturn(null); + VeniceException ex7 = expectThrows( + VeniceException.class, + () -> createVersionOk.handleStreamPushType(admin, store, request, response, Lazy.of(() -> true))); + assertTrue(ex7.getMessage().contains("No hybrid version found"), "Got: " + ex7.getMessage()); + + // Case 4: Child region; Non-aggregate mode and there is a hybrid version + Version mockVersion = mock(Version.class); + when(mockVersion.getPartitionCount()).thenReturn(42); + when(admin.isParent()).thenReturn(false); + when(store.getHybridStoreConfig().getDataReplicationPolicy()).thenReturn(DataReplicationPolicy.NON_AGGREGATE); + when(admin.getReferenceVersionForStreamingWrites(anyString(), anyString(), anyString())).thenReturn(mockVersion); + createVersionOk.handleStreamPushType(admin, store, request, response, Lazy.of(() -> true)); + assertEquals(response.getPartitions(), 42); + assertEquals(response.getCompressionStrategy(), CompressionStrategy.NO_OP); + assertEquals(response.getKafkaTopic(), Version.composeRealTimeTopic(STORE_NAME)); + } + + @Test + public void testGetActiveActiveReplicationCheck() { + Admin admin = mock(Admin.class); + Store store = mock(Store.class); + String clusterName = "testCluster"; + String storeName = "testStore"; + CreateVersion createVersion = new CreateVersion(true, Optional.of(accessClient), false, false); + + // Case 1: Admin is parent, store has AA replication, and AA replication is enabled in all regions + when(admin.isParent()).thenReturn(true); + when(store.isActiveActiveReplicationEnabled()).thenReturn(true); + when(admin.isActiveActiveReplicationEnabledInAllRegion(clusterName, storeName, true)).thenReturn(true); + + Lazy check = createVersion.getActiveActiveReplicationCheck(admin, store, clusterName, storeName, true); + assertTrue(check.get(), "Expected AA replication check to return true"); + + // Case 2: Admin is not parent + when(admin.isParent()).thenReturn(false); + check = createVersion.getActiveActiveReplicationCheck(admin, store, clusterName, storeName, true); + assertFalse(check.get(), "Expected AA replication check to return false as admin is not parent"); + + // Case 3: Store does not have AA replication enabled + when(admin.isParent()).thenReturn(true); + when(store.isActiveActiveReplicationEnabled()).thenReturn(false); + check = createVersion.getActiveActiveReplicationCheck(admin, store, clusterName, storeName, true); + assertFalse(check.get(), "Expected AA replication check to return false as store does not have AA replication"); + } + + @Test + public void testApplyConfigBasedOnReplication() { + Lazy isAARCheckEnabled = Lazy.of(() -> true); + String configType = "TestConfig"; + String configValue = "TestValue"; + String storeName = "testStore"; + + CreateVersion createVersion = new CreateVersion(true, Optional.of(accessClient), false, false); + + // Case 1: Config is applied as AA replication is enabled + String result = createVersion.applyConfigBasedOnReplication(configType, configValue, storeName, isAARCheckEnabled); + assertEquals(result, configValue, "Expected config to be applied as AA replication is enabled"); + + // Case 2: Config is ignored as AA replication is disabled + isAARCheckEnabled = Lazy.of(() -> false); + result = createVersion.applyConfigBasedOnReplication(configType, configValue, storeName, isAARCheckEnabled); + assertNull(result, "Expected config to be ignored as AA replication is disabled"); + + // Case 3: Config value is null + result = createVersion.applyConfigBasedOnReplication(configType, null, storeName, isAARCheckEnabled); + assertNull(result, "Expected config to remain null when input configValue is null"); + } + + @Test + public void testHandleNonStreamPushType() { + String clusterName = "testCluster"; + String storeName = "testStore"; + String pushJobId = "pushJob123"; + int versionNumber = 11; + Version.PushType pushType = INCREMENTAL; + int computedPartitionCount = 10; + Admin admin = mock(Admin.class); + Store store = mock(Store.class); + RequestTopicForPushRequest request = new RequestTopicForPushRequest(clusterName, storeName, pushType, pushJobId); + VersionCreationResponse response = new VersionCreationResponse(); + CreateVersion createVersion = new CreateVersion(true, Optional.of(accessClient), false, false); + Lazy isActiveActiveReplicationEnabledInAllRegions = Lazy.of(() -> true); + + // Mock admin methods + when(admin.whetherEnableBatchPushFromAdmin(storeName)).thenReturn(true); + when(admin.calculateNumberOfPartitions(clusterName, storeName)).thenReturn(computedPartitionCount); + + Version version = mock(Version.class); + when(version.getStoreName()).thenReturn(storeName); + when(version.getPartitionCount()).thenReturn(computedPartitionCount); + when(version.getNumber()).thenReturn(versionNumber); + + when( + admin.incrementVersionIdempotent( + clusterName, + storeName, + request.getPushJobId(), + computedPartitionCount, + response.getReplicas(), + pushType, + request.isSendStartOfPush(), + request.isSorted(), + request.getCompressionDictionary(), + Optional.ofNullable(request.getSourceGridFabric()), + Optional.ofNullable(request.getCertificateInRequest()), + request.getRewindTimeInSecondsOverride(), + Optional.ofNullable(request.getEmergencySourceRegion()), + request.isDeferVersionSwap(), + request.getTargetedRegions(), + request.getRepushSourceVersion())).thenReturn(version); + + when(createVersion.getCompressionStrategy(version, "testStore_v1")).thenReturn(CompressionStrategy.NO_OP); + + // Case 1: Happy Path - All validations pass + createVersion + .handleNonStreamPushType(admin, store, request, response, isActiveActiveReplicationEnabledInAllRegions); + assertEquals(response.getPartitions(), computedPartitionCount, "Expected partition count to match."); + assertEquals(response.getVersion(), versionNumber, "Expected version number to match."); + assertEquals(response.getKafkaTopic(), "testStore_rt", "Expected Kafka topic to match."); + assertEquals( + response.getCompressionStrategy(), + CompressionStrategy.NO_OP, + "Expected compression strategy to be NO_OP."); + + // Case 2: Batch push is not enabled + when(admin.whetherEnableBatchPushFromAdmin(storeName)).thenReturn(false); + VeniceUnsupportedOperationException ex1 = expectThrows( + VeniceUnsupportedOperationException.class, + () -> createVersion + .handleNonStreamPushType(admin, store, request, response, isActiveActiveReplicationEnabledInAllRegions)); + assertTrue(ex1.getMessage().contains("Please push data to Venice Parent Colo instead")); + + // Case 3: Increment version fails + doThrow(new VeniceException("Version creation failure")).when(admin) + .incrementVersionIdempotent( + clusterName, + storeName, + request.getPushJobId(), + computedPartitionCount, + response.getReplicas(), + pushType, + request.isSendStartOfPush(), + request.isSorted(), + request.getCompressionDictionary(), + Optional.ofNullable(request.getSourceGridFabric()), + Optional.ofNullable(request.getCertificateInRequest()), + request.getRewindTimeInSecondsOverride(), + Optional.ofNullable(request.getEmergencySourceRegion()), + request.isDeferVersionSwap(), + request.getTargetedRegions(), + request.getRepushSourceVersion()); + + when(admin.whetherEnableBatchPushFromAdmin(storeName)).thenReturn(true); + VeniceException ex2 = expectThrows( + VeniceException.class, + () -> createVersion + .handleNonStreamPushType(admin, store, request, response, isActiveActiveReplicationEnabledInAllRegions)); + assertTrue(ex2.getMessage().contains("Version creation failure"), "Actual Message: " + ex2.getMessage()); + } } diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/TestVeniceRouteHandler.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/TestVeniceRouteHandler.java index f01b0650572..690d2d8cd97 100644 --- a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/TestVeniceRouteHandler.java +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/TestVeniceRouteHandler.java @@ -1,15 +1,20 @@ package com.linkedin.venice.controller.server; +import static org.mockito.Mockito.doReturn; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; +import com.linkedin.venice.controller.Admin; +import com.linkedin.venice.controller.ControllerRequestHandlerDependencies; +import com.linkedin.venice.controller.VeniceParentHelixAdmin; import com.linkedin.venice.controllerapi.ControllerResponse; import com.linkedin.venice.exceptions.ErrorType; import com.linkedin.venice.exceptions.ExceptionType; import com.linkedin.venice.utils.ObjectMapperFactory; import org.apache.commons.httpclient.HttpStatus; import org.testng.Assert; +import org.testng.annotations.BeforeMethod; import org.testng.annotations.Test; import spark.Request; import spark.Response; @@ -17,6 +22,15 @@ public class TestVeniceRouteHandler { + private Admin mockAdmin; + + @BeforeMethod + public void setUp() { + mockAdmin = mock(VeniceParentHelixAdmin.class); + ControllerRequestHandlerDependencies dependencies = mock(ControllerRequestHandlerDependencies.class); + doReturn(mockAdmin).when(dependencies).getAdmin(); + } + @Test public void testIsAllowListUser() throws Exception { Route userAllowedRoute = new VeniceRouteHandler(ControllerResponse.class) { diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/VeniceControllerAccessManagerTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/VeniceControllerAccessManagerTest.java new file mode 100644 index 00000000000..8084035eb6d --- /dev/null +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/VeniceControllerAccessManagerTest.java @@ -0,0 +1,132 @@ +package com.linkedin.venice.controller.server; + +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import com.linkedin.venice.acl.AclException; +import com.linkedin.venice.acl.DynamicAccessController; +import com.linkedin.venice.acl.NoOpDynamicAccessController; +import com.linkedin.venice.authorization.Method; +import java.security.cert.X509Certificate; +import javax.security.auth.x500.X500Principal; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class VeniceControllerAccessManagerTest { + private DynamicAccessController mockedAccessController; + private VeniceControllerAccessManager controllerAccessManager; + + @BeforeMethod + public void setUp() { + mockedAccessController = mock(DynamicAccessController.class); + controllerAccessManager = new VeniceControllerAccessManager(mockedAccessController); + } + + @Test + public void testHasWriteAccessToPubSubTopicGranted() throws AclException { + String resourceName = "test-topic"; + X509Certificate mockCert = mock(X509Certificate.class); + + when(mockedAccessController.hasAccess(mockCert, resourceName, Method.Write.name())).thenReturn(true); + + assertTrue(controllerAccessManager.hasWriteAccessToPubSubTopic(resourceName, mockCert, "host", "127.0.0.1")); + verify(mockedAccessController, times(1)).hasAccess(mockCert, resourceName, Method.Write.name()); + } + + @Test + public void testHasWriteAccessToPubSubTopicDenied() throws AclException { + String resourceName = "test-topic"; + X509Certificate mockCert = mock(X509Certificate.class); + + when(mockedAccessController.hasAccess(mockCert, resourceName, Method.Write.name())).thenReturn(false); + + assertFalse(controllerAccessManager.hasWriteAccessToPubSubTopic(resourceName, mockCert, "host", "127.0.0.1")); + verify(mockedAccessController, times(1)).hasAccess(mockCert, resourceName, Method.Write.name()); + } + + @Test + public void testHasAccessHandlesNullCertificate() throws AclException { + String resourceName = "test-topic"; + + when(mockedAccessController.hasAccess(null, resourceName, Method.Read.name())).thenReturn(true); + + assertTrue(controllerAccessManager.hasReadAccessToPubSubTopic(resourceName, null, "host", "127.0.0.1")); + verify(mockedAccessController, times(1)).hasAccess(null, resourceName, Method.Read.name()); + } + + @Test + public void testHasAccessHandlesNullResource() { + X509Certificate mockCert = mock(X509Certificate.class); + + Exception e = expectThrows( + NullPointerException.class, + () -> controllerAccessManager.hasAccess(null, mockCert, Method.Write, "host", "")); + assertTrue(e.getMessage().contains("Resource name is required to enforce ACL")); + } + + @Test + public void testHasAccessHandlesNullMethod() { + String resourceName = "test-topic"; + X509Certificate mockCert = mock(X509Certificate.class); + + Exception e = expectThrows( + NullPointerException.class, + () -> controllerAccessManager.hasAccess(resourceName, mockCert, null, "host", "127.0.0.1")); + assertTrue(e.getMessage().contains("Access method is required to enforce ACL")); + } + + @Test + public void testHasAccessWhenAccessControllerThrowsException() throws AclException { + String resourceName = "test-topic"; + X509Certificate mockCert = mock(X509Certificate.class); + + when(mockedAccessController.hasAccess(mockCert, resourceName, Method.Write.name())) + .thenThrow(new AclException("Test exception")); + + assertFalse(controllerAccessManager.hasAccess(resourceName, mockCert, Method.Write, "host", "")); + verify(mockedAccessController, times(1)).hasAccess(mockCert, resourceName, Method.Write.name()); + } + + @Test + public void testHasAccessWhenAccessControllerIsNull() { + Exception e = expectThrows(NullPointerException.class, () -> new VeniceControllerAccessManager(null)); + assertTrue(e.getMessage().contains("is required to enforce ACL")); + } + + @Test + public void testAccessManagerWithNoOpDynamicAccessController() { + String resourceName = "test-topic"; + String hostname = "host"; + String remoteAddr = "127.0.0.1"; + + // Case 1: Non-null certificate + X509Certificate mockCert = mock(X509Certificate.class); + X500Principal principal = new X500Principal("CN=Test User"); + doReturn(principal).when(mockCert).getSubjectX500Principal(); + + VeniceControllerAccessManager accessManager = + new VeniceControllerAccessManager(NoOpDynamicAccessController.INSTANCE); + assertTrue(accessManager.hasAccessToStore(resourceName, mockCert, hostname, remoteAddr)); + assertTrue(accessManager.hasReadAccessToPubSubTopic(resourceName, mockCert, hostname, remoteAddr)); + assertTrue(accessManager.hasWriteAccessToPubSubTopic(resourceName, mockCert, hostname, remoteAddr)); + assertTrue(accessManager.isAllowListUser(resourceName, mockCert)); + assertFalse(accessManager.isAclEnabled()); + assertEquals(accessManager.getPrincipalId(mockCert), "CN=Test User"); + + // Case 2: Null certificate + assertTrue(accessManager.hasAccessToStore(resourceName, null, hostname, remoteAddr)); + assertTrue(accessManager.hasReadAccessToPubSubTopic(resourceName, null, hostname, remoteAddr)); + assertTrue(accessManager.hasWriteAccessToPubSubTopic(resourceName, null, hostname, remoteAddr)); + assertTrue(accessManager.isAllowListUser(resourceName, null)); + assertFalse(accessManager.isAclEnabled()); + assertEquals(accessManager.getPrincipalId(null), VeniceControllerAccessManager.UNKNOWN_USER); + } +} diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/VeniceControllerGrpcServiceImplTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/VeniceControllerGrpcServiceImplTest.java new file mode 100644 index 00000000000..985046fcdd6 --- /dev/null +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/VeniceControllerGrpcServiceImplTest.java @@ -0,0 +1,236 @@ +package com.linkedin.venice.controller.server; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.mock; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; +import static org.testng.Assert.expectThrows; + +import com.linkedin.venice.controllerapi.transport.GrpcRequestResponseConverter; +import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo; +import com.linkedin.venice.protocols.controller.ControllerGrpcErrorType; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcRequest; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcResponse; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcRequest; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcResponse; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcRequest; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcResponse; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcErrorInfo; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcServiceGrpc; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcServiceGrpc.VeniceControllerGrpcServiceBlockingStub; +import io.grpc.ManagedChannel; +import io.grpc.Server; +import io.grpc.Status; +import io.grpc.StatusRuntimeException; +import io.grpc.inprocess.InProcessChannelBuilder; +import io.grpc.inprocess.InProcessServerBuilder; +import org.testng.annotations.AfterMethod; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class VeniceControllerGrpcServiceImplTest { + private static final String TEST_CLUSTER = "test-cluster"; + private static final String TEST_STORE = "test-store"; + private static final String D2_TEST_SERVICE = "d2://test-service"; + private static final String D2_TEST_SERVER = "d2://test-server"; + private static final String HTTP_URL = "http://localhost:8080"; + private static final String HTTPS_URL = "https://localhost:8081"; + private static final String GRPC_URL = "grpc://localhost:8082"; + private static final String SECURE_GRPC_URL = "grpcs://localhost:8083"; + private static final String OWNER = "test-owner"; + private static final String KEY_SCHEMA = "int"; + private static final String VALUE_SCHEMA = "string"; + + private Server grpcServer; + private ManagedChannel grpcChannel; + private VeniceControllerRequestHandler requestHandler; + private VeniceControllerGrpcServiceBlockingStub blockingStub; + + @BeforeMethod + public void setUp() throws Exception { + requestHandler = mock(VeniceControllerRequestHandler.class); + + // Create a unique server name for the in-process server + String serverName = InProcessServerBuilder.generateName(); + + // Start the gRPC server in-process + grpcServer = InProcessServerBuilder.forName(serverName) + .directExecutor() + .addService(new VeniceControllerGrpcServiceImpl(requestHandler)) + .build() + .start(); + + // Create a channel to communicate with the server + grpcChannel = InProcessChannelBuilder.forName(serverName).directExecutor().build(); + + // Create a blocking stub to make calls to the server + blockingStub = VeniceControllerGrpcServiceGrpc.newBlockingStub(grpcChannel); + } + + @AfterMethod + public void tearDown() throws Exception { + if (grpcServer != null) { + grpcServer.shutdown(); + } + if (grpcChannel != null) { + grpcChannel.shutdown(); + } + } + + @Test + public void testGetLeaderController() { + // Case 1: Successful response + LeaderControllerGrpcResponse response = LeaderControllerGrpcResponse.newBuilder() + .setClusterName(TEST_CLUSTER) + .setHttpUrl(HTTP_URL) + .setHttpsUrl(HTTPS_URL) + .setGrpcUrl(GRPC_URL) + .setSecureGrpcUrl(SECURE_GRPC_URL) + .build(); + doReturn(response).when(requestHandler).getLeaderControllerDetails(any(LeaderControllerGrpcRequest.class)); + + LeaderControllerGrpcRequest request = LeaderControllerGrpcRequest.newBuilder().setClusterName(TEST_CLUSTER).build(); + LeaderControllerGrpcResponse actualResponse = blockingStub.getLeaderController(request); + + assertNotNull(actualResponse, "Response should not be null"); + assertEquals(actualResponse.getClusterName(), TEST_CLUSTER, "Cluster name should match"); + assertEquals(actualResponse.getHttpUrl(), HTTP_URL, "HTTP URL should match"); + assertEquals(actualResponse.getHttpsUrl(), HTTPS_URL, "HTTPS URL should match"); + assertEquals(actualResponse.getGrpcUrl(), GRPC_URL, "gRPC URL should match"); + assertEquals(actualResponse.getSecureGrpcUrl(), SECURE_GRPC_URL, "Secure gRPC URL should match"); + + // Case 2: Bad request as cluster name is missing + doThrow(new IllegalArgumentException("Cluster name is required for leader controller discovery")) + .when(requestHandler) + .getLeaderControllerDetails(any(LeaderControllerGrpcRequest.class)); + LeaderControllerGrpcRequest requestWithoutClusterName = LeaderControllerGrpcRequest.newBuilder().build(); + StatusRuntimeException e = + expectThrows(StatusRuntimeException.class, () -> blockingStub.getLeaderController(requestWithoutClusterName)); + assertNotNull(e.getStatus(), "Status should not be null"); + assertEquals(e.getStatus().getCode(), Status.INVALID_ARGUMENT.getCode()); + + VeniceControllerGrpcErrorInfo errorInfo = GrpcRequestResponseConverter.parseControllerGrpcError(e); + assertNotNull(errorInfo, "Error info should not be null"); + assertFalse(errorInfo.hasStoreName(), "Store name should not be present in the error info"); + assertEquals(errorInfo.getErrorType(), ControllerGrpcErrorType.BAD_REQUEST); + assertTrue(errorInfo.getErrorMessage().contains("Cluster name is required for leader controller discovery")); + + // Case 3: requestHandler throws an exception + doThrow(new VeniceException("Failed to get leader controller")).when(requestHandler) + .getLeaderControllerDetails(any(LeaderControllerGrpcRequest.class)); + StatusRuntimeException e2 = + expectThrows(StatusRuntimeException.class, () -> blockingStub.getLeaderController(request)); + assertNotNull(e2.getStatus(), "Status should not be null"); + assertEquals(e2.getStatus().getCode(), Status.INTERNAL.getCode()); + VeniceControllerGrpcErrorInfo errorInfo2 = GrpcRequestResponseConverter.parseControllerGrpcError(e2); + assertNotNull(errorInfo2, "Error info should not be null"); + assertTrue(errorInfo2.hasClusterName(), "Cluster name should be present in the error info"); + assertEquals(errorInfo2.getClusterName(), TEST_CLUSTER); + assertFalse(errorInfo2.hasStoreName(), "Store name should not be present in the error info"); + assertEquals(errorInfo2.getErrorType(), ControllerGrpcErrorType.GENERAL_ERROR); + assertTrue(errorInfo2.getErrorMessage().contains("Failed to get leader controller")); + } + + @Test + public void testDiscoverClusterForStore() { + // Case 1: Successful response + DiscoverClusterGrpcResponse response = DiscoverClusterGrpcResponse.newBuilder() + .setStoreName(TEST_STORE) + .setClusterName(TEST_CLUSTER) + .setD2Service(D2_TEST_SERVICE) + .setServerD2Service(D2_TEST_SERVER) + .build(); + doReturn(response).when(requestHandler).discoverCluster(any(DiscoverClusterGrpcRequest.class)); + DiscoverClusterGrpcRequest request = DiscoverClusterGrpcRequest.newBuilder().setStoreName(TEST_STORE).build(); + DiscoverClusterGrpcResponse actualResponse = blockingStub.discoverClusterForStore(request); + assertNotNull(actualResponse, "Response should not be null"); + assertEquals(actualResponse.getStoreName(), TEST_STORE, "Store name should match"); + assertEquals(actualResponse.getClusterName(), TEST_CLUSTER, "Cluster name should match"); + assertEquals(actualResponse.getD2Service(), D2_TEST_SERVICE, "D2 service should match"); + assertEquals(actualResponse.getServerD2Service(), D2_TEST_SERVER, "Server D2 service should match"); + + // Case 2: Bad request as store name is missing + doThrow(new IllegalArgumentException("Store name is required for cluster discovery")).when(requestHandler) + .discoverCluster(any(DiscoverClusterGrpcRequest.class)); + DiscoverClusterGrpcRequest requestWithoutStoreName = DiscoverClusterGrpcRequest.newBuilder().build(); + StatusRuntimeException e = + expectThrows(StatusRuntimeException.class, () -> blockingStub.discoverClusterForStore(requestWithoutStoreName)); + assertNotNull(e.getStatus(), "Status should not be null"); + assertEquals(e.getStatus().getCode(), Status.INVALID_ARGUMENT.getCode()); + VeniceControllerGrpcErrorInfo errorInfo = GrpcRequestResponseConverter.parseControllerGrpcError(e); + assertNotNull(errorInfo, "Error info should not be null"); + assertFalse(errorInfo.hasClusterName(), "Cluster name should not be present in the error info"); + assertEquals(errorInfo.getErrorType(), ControllerGrpcErrorType.BAD_REQUEST); + assertTrue(errorInfo.getErrorMessage().contains("Store name is required for cluster discovery")); + + // Case 3: requestHandler throws an exception + doThrow(new VeniceException("Failed to discover cluster")).when(requestHandler) + .discoverCluster(any(DiscoverClusterGrpcRequest.class)); + StatusRuntimeException e2 = + expectThrows(StatusRuntimeException.class, () -> blockingStub.discoverClusterForStore(request)); + assertNotNull(e2.getStatus(), "Status should not be null"); + assertEquals(e2.getStatus().getCode(), Status.INTERNAL.getCode()); + VeniceControllerGrpcErrorInfo errorInfo2 = GrpcRequestResponseConverter.parseControllerGrpcError(e2); + assertNotNull(errorInfo2, "Error info should not be null"); + assertFalse(errorInfo2.hasClusterName(), "Cluster name should not be present in the error info"); + assertEquals(errorInfo2.getErrorType(), ControllerGrpcErrorType.GENERAL_ERROR); + assertTrue(errorInfo2.getErrorMessage().contains("Failed to discover cluster")); + } + + @Test + public void testCreateStore() { + CreateStoreGrpcResponse response = CreateStoreGrpcResponse.newBuilder() + .setClusterStoreInfo( + ClusterStoreGrpcInfo.newBuilder().setClusterName(TEST_CLUSTER).setStoreName(TEST_STORE).build()) + .setOwner(OWNER) + .build(); + // Case 1: Successful response + doReturn(response).when(requestHandler).createStore(any(CreateStoreGrpcRequest.class)); + CreateStoreGrpcRequest request = CreateStoreGrpcRequest.newBuilder() + .setClusterStoreInfo( + ClusterStoreGrpcInfo.newBuilder().setClusterName(TEST_CLUSTER).setStoreName(TEST_STORE).build()) + .setOwner(OWNER) + .setKeySchema(KEY_SCHEMA) + .setValueSchema(VALUE_SCHEMA) + .build(); + CreateStoreGrpcResponse actualResponse = blockingStub.createStore(request); + assertNotNull(actualResponse, "Response should not be null"); + assertNotNull(actualResponse.getClusterStoreInfo(), "ClusterStoreInfo should not be null"); + assertEquals(actualResponse.getClusterStoreInfo().getClusterName(), TEST_CLUSTER, "Cluster name should match"); + assertEquals(actualResponse.getClusterStoreInfo().getStoreName(), TEST_STORE, "Store name should match"); + + // Case 2: Bad request as cluster name is missing + CreateStoreGrpcRequest requestWithoutClusterName = CreateStoreGrpcRequest.newBuilder() + .setOwner(OWNER) + .setKeySchema(KEY_SCHEMA) + .setValueSchema(VALUE_SCHEMA) + .build(); + doThrow(new IllegalArgumentException("The request is missing the cluster_name")).when(requestHandler) + .createStore(any(CreateStoreGrpcRequest.class)); + StatusRuntimeException e = + expectThrows(StatusRuntimeException.class, () -> blockingStub.createStore(requestWithoutClusterName)); + assertNotNull(e.getStatus(), "Status should not be null"); + assertEquals(e.getStatus().getCode(), Status.INVALID_ARGUMENT.getCode()); + VeniceControllerGrpcErrorInfo errorInfo = GrpcRequestResponseConverter.parseControllerGrpcError(e); + assertEquals(errorInfo.getErrorType(), ControllerGrpcErrorType.BAD_REQUEST); + assertNotNull(errorInfo, "Error info should not be null"); + assertTrue(errorInfo.getErrorMessage().contains("The request is missing the cluster_name")); + + // Case 3: requestHandler throws an exception + doThrow(new VeniceException("Failed to create store")).when(requestHandler) + .createStore(any(CreateStoreGrpcRequest.class)); + StatusRuntimeException e3 = expectThrows(StatusRuntimeException.class, () -> blockingStub.createStore(request)); + assertNotNull(e3.getStatus(), "Status should not be null"); + assertEquals(e3.getStatus().getCode(), Status.INTERNAL.getCode()); + VeniceControllerGrpcErrorInfo errorInfo3 = GrpcRequestResponseConverter.parseControllerGrpcError(e3); + assertNotNull(errorInfo3, "Error info should not be null"); + assertEquals(errorInfo3.getErrorType(), ControllerGrpcErrorType.GENERAL_ERROR); + assertTrue(errorInfo3.getErrorMessage().contains("Failed to create store")); + } +} diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/VeniceControllerRequestHandlerTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/VeniceControllerRequestHandlerTest.java new file mode 100644 index 00000000000..4fca6d482a3 --- /dev/null +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/VeniceControllerRequestHandlerTest.java @@ -0,0 +1,147 @@ +package com.linkedin.venice.controller.server; + +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + +import com.linkedin.venice.controller.Admin; +import com.linkedin.venice.controller.ControllerRequestHandlerDependencies; +import com.linkedin.venice.meta.Instance; +import com.linkedin.venice.protocols.controller.ClusterStoreGrpcInfo; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcRequest; +import com.linkedin.venice.protocols.controller.CreateStoreGrpcResponse; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcRequest; +import com.linkedin.venice.protocols.controller.DiscoverClusterGrpcResponse; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcRequest; +import com.linkedin.venice.protocols.controller.LeaderControllerGrpcResponse; +import com.linkedin.venice.utils.Pair; +import java.util.Optional; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class VeniceControllerRequestHandlerTest { + private VeniceControllerRequestHandler requestHandler; + private Admin admin; + private ControllerRequestHandlerDependencies dependencies; + + @BeforeMethod + public void setUp() { + admin = mock(Admin.class); + dependencies = mock(ControllerRequestHandlerDependencies.class); + when(dependencies.getAdmin()).thenReturn(admin); + when(dependencies.isSslEnabled()).thenReturn(true); + requestHandler = new VeniceControllerRequestHandler(dependencies); + } + + @Test + public void testGetLeaderControllerDetails() { + String clusterName = "testCluster"; + LeaderControllerGrpcRequest request = LeaderControllerGrpcRequest.newBuilder().setClusterName(clusterName).build(); + Instance leaderInstance = mock(Instance.class); + when(admin.getLeaderController(clusterName)).thenReturn(leaderInstance); + when(leaderInstance.getUrl(true)).thenReturn("https://leader-url:443"); + when(leaderInstance.getUrl(false)).thenReturn("http://leader-url:80"); + when(leaderInstance.getGrpcUrl()).thenReturn("leader-grpc-url:50051"); + when(leaderInstance.getGrpcSslUrl()).thenReturn("leader-grpc-url:50052"); + when(leaderInstance.getPort()).thenReturn(80); + when(leaderInstance.getSslPort()).thenReturn(443); // SSL enabled + when(leaderInstance.getGrpcPort()).thenReturn(50051); + when(leaderInstance.getGrpcSslPort()).thenReturn(50052); + + LeaderControllerGrpcResponse response = requestHandler.getLeaderControllerDetails(request); + + assertEquals(response.getClusterName(), clusterName); + assertEquals(response.getHttpUrl(), "https://leader-url:443"); // SSL enabled + assertEquals(response.getHttpsUrl(), "https://leader-url:443"); + assertEquals(response.getGrpcUrl(), "leader-grpc-url:50051"); + assertEquals(response.getSecureGrpcUrl(), "leader-grpc-url:50052"); + + // SSL not enabled + when(dependencies.isSslEnabled()).thenReturn(false); + requestHandler = new VeniceControllerRequestHandler(dependencies); + LeaderControllerGrpcResponse response1 = requestHandler.getLeaderControllerDetails(request); + assertEquals(response1.getHttpUrl(), "http://leader-url:80"); + assertEquals(response1.getHttpsUrl(), "https://leader-url:443"); + assertEquals(response1.getGrpcUrl(), "leader-grpc-url:50051"); + assertEquals(response1.getSecureGrpcUrl(), "leader-grpc-url:50052"); + } + + @Test + public void testDiscoverCluster() { + String storeName = "testStore"; + DiscoverClusterGrpcRequest request = DiscoverClusterGrpcRequest.newBuilder().setStoreName(storeName).build(); + Pair clusterToD2Pair = Pair.create("testCluster", "testD2Service"); + when(admin.discoverCluster(storeName)).thenReturn(clusterToD2Pair); + when(admin.getServerD2Service("testCluster")).thenReturn("testServerD2Service"); + + DiscoverClusterGrpcResponse response = requestHandler.discoverCluster(request); + + assertEquals(response.getStoreName(), storeName); + assertEquals(response.getClusterName(), "testCluster"); + assertEquals(response.getD2Service(), "testD2Service"); + assertEquals(response.getServerD2Service(), "testServerD2Service"); + } + + @Test + public void testCreateStore() { + CreateStoreGrpcRequest request = CreateStoreGrpcRequest.newBuilder() + .setClusterStoreInfo( + ClusterStoreGrpcInfo.newBuilder().setClusterName("testCluster").setStoreName("testStore").build()) + .setKeySchema("testKeySchema") + .setValueSchema("testValueSchema") + .setOwner("testOwner") + .setAccessPermission("testAccessPermissions") + .setIsSystemStore(false) + .build(); + + CreateStoreGrpcResponse response = requestHandler.createStore(request); + + verify(admin, times(1)).createStore( + "testCluster", + "testStore", + "testOwner", + "testKeySchema", + "testValueSchema", + false, + Optional.of("testAccessPermissions")); + assertEquals(response.getClusterStoreInfo().getClusterName(), "testCluster"); + assertEquals(response.getClusterStoreInfo().getStoreName(), "testStore"); + assertEquals(response.getOwner(), "testOwner"); + } + + @Test + public void testCreateStoreWithNullAccessPermissions() { + CreateStoreGrpcRequest request = CreateStoreGrpcRequest.newBuilder() + .setClusterStoreInfo( + ClusterStoreGrpcInfo.newBuilder().setClusterName("testCluster").setStoreName("testStore").build()) + .setKeySchema("testKeySchema") + .setValueSchema("testValueSchema") + .setOwner("testOwner") + .setIsSystemStore(true) + .build(); + + CreateStoreGrpcResponse response = requestHandler.createStore(request); + + verify(admin, times(1)).createStore( + "testCluster", + "testStore", + "testOwner", + "testKeySchema", + "testValueSchema", + true, + Optional.empty()); + assertEquals(response.getClusterStoreInfo().getClusterName(), "testCluster"); + assertEquals(response.getClusterStoreInfo().getStoreName(), "testStore"); + assertEquals(response.getOwner(), "testOwner"); + } + + @Test + public void testIsSslEnabled() { + boolean sslEnabled = requestHandler.isSslEnabled(); + assertTrue(sslEnabled); + } +} diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/grpc/ControllerGrpcSslSessionInterceptorTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/grpc/ControllerGrpcSslSessionInterceptorTest.java new file mode 100644 index 00000000000..47eb6b34070 --- /dev/null +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/grpc/ControllerGrpcSslSessionInterceptorTest.java @@ -0,0 +1,108 @@ +package com.linkedin.venice.controller.server.grpc; + +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import io.grpc.Attributes; +import io.grpc.Grpc; +import io.grpc.Metadata; +import io.grpc.ServerCall; +import io.grpc.ServerCallHandler; +import java.net.SocketAddress; +import java.security.cert.Certificate; +import java.security.cert.X509Certificate; +import javax.net.ssl.SSLPeerUnverifiedException; +import javax.net.ssl.SSLSession; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class ControllerGrpcSslSessionInterceptorTest { + private ControllerGrpcSslSessionInterceptor interceptor; + private ServerCall serverCall; + private Metadata metadata; + private ServerCallHandler serverCallHandler; + + @BeforeMethod + public void setUp() { + interceptor = new ControllerGrpcSslSessionInterceptor(); + serverCall = mock(ServerCall.class); + metadata = new Metadata(); + serverCallHandler = mock(ServerCallHandler.class); + } + + @Test + public void testSslSessionNotEnabled() { + Attributes attributes = Attributes.newBuilder().build(); + when(serverCall.getAttributes()).thenReturn(attributes); + interceptor.interceptCall(serverCall, metadata, serverCallHandler); + verify(serverCall, times(1)).close( + eq(ControllerGrpcSslSessionInterceptor.NON_SSL_CONNECTION_ERROR.getStatus()), + eq(ControllerGrpcSslSessionInterceptor.NON_SSL_CONNECTION_ERROR.getTrailers())); + verify(serverCallHandler, never()).startCall(any(), any()); + } + + @Test + public void testSslSessionEnabledWithValidCertificate() throws SSLPeerUnverifiedException { + SSLSession sslSession = mock(SSLSession.class); + SocketAddress remoteAddress = mock(SocketAddress.class); + X509Certificate clientCertificate = mock(X509Certificate.class); + + Attributes attributes = Attributes.newBuilder() + .set(Grpc.TRANSPORT_ATTR_SSL_SESSION, sslSession) + .set(Grpc.TRANSPORT_ATTR_REMOTE_ADDR, remoteAddress) + .build(); + + when(serverCall.getAttributes()).thenReturn(attributes); + when(remoteAddress.toString()).thenReturn("remote-address"); + Certificate[] peerCertificates = new Certificate[] { clientCertificate }; + when(sslSession.getPeerCertificates()).thenReturn(peerCertificates); + + interceptor.interceptCall(serverCall, metadata, serverCallHandler); + + verify(serverCallHandler, times(1)).startCall(serverCall, metadata); + } + + @Test + public void testCertificateExtractionFails() throws SSLPeerUnverifiedException { + SSLSession sslSession = mock(SSLSession.class); + SocketAddress remoteAddress = mock(SocketAddress.class); + + Attributes attributes = Attributes.newBuilder() + .set(Grpc.TRANSPORT_ATTR_SSL_SESSION, sslSession) + .set(Grpc.TRANSPORT_ATTR_REMOTE_ADDR, remoteAddress) + .build(); + + when(serverCall.getAttributes()).thenReturn(attributes); + doThrow(new RuntimeException("Failed to extract certificate")).when(sslSession).getPeerCertificates(); + + interceptor.interceptCall(serverCall, metadata, serverCallHandler); + + verify(serverCall, times(1)).close( + eq(ControllerGrpcSslSessionInterceptor.NON_SSL_CONNECTION_ERROR.getStatus()), + eq(ControllerGrpcSslSessionInterceptor.NON_SSL_CONNECTION_ERROR.getTrailers())); + verify(serverCallHandler, never()).startCall(any(), any()); + } + + @Test + public void testRemoteAddressUnknown() throws SSLPeerUnverifiedException { + SSLSession sslSession = mock(SSLSession.class); + X509Certificate clientCertificate = mock(X509Certificate.class); + + Attributes attributes = Attributes.newBuilder().set(Grpc.TRANSPORT_ATTR_SSL_SESSION, sslSession).build(); + + when(serverCall.getAttributes()).thenReturn(attributes); + Certificate[] peerCertificates = new Certificate[] { clientCertificate }; + when(sslSession.getPeerCertificates()).thenReturn(peerCertificates); + + interceptor.interceptCall(serverCall, metadata, serverCallHandler); + + verify(serverCallHandler, times(1)).startCall(serverCall, metadata); + } +} diff --git a/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/grpc/ParentControllerRegionValidationInterceptorTest.java b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/grpc/ParentControllerRegionValidationInterceptorTest.java new file mode 100644 index 00000000000..903be8bf835 --- /dev/null +++ b/services/venice-controller/src/test/java/com/linkedin/venice/controller/server/grpc/ParentControllerRegionValidationInterceptorTest.java @@ -0,0 +1,97 @@ +package com.linkedin.venice.controller.server.grpc; + +import static org.mockito.Mockito.RETURNS_DEEP_STUBS; +import static org.mockito.Mockito.any; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + +import com.linkedin.venice.controller.Admin; +import com.linkedin.venice.controller.ParentControllerRegionState; +import com.linkedin.venice.protocols.controller.VeniceControllerGrpcServiceGrpc; +import io.grpc.Attributes; +import io.grpc.Metadata; +import io.grpc.MethodDescriptor; +import io.grpc.ServerCall; +import io.grpc.ServerCallHandler; +import org.mockito.ArgumentCaptor; +import org.testng.annotations.BeforeMethod; +import org.testng.annotations.Test; + + +public class ParentControllerRegionValidationInterceptorTest { + private ParentControllerRegionValidationInterceptor interceptor; + private Admin admin; + private ServerCall call; + private Metadata headers; + private ServerCallHandler next; + + @BeforeMethod + public void setUp() { + admin = mock(Admin.class); + call = mock(ServerCall.class, RETURNS_DEEP_STUBS); + headers = new Metadata(); + next = mock(ServerCallHandler.class); + interceptor = new ParentControllerRegionValidationInterceptor(admin); + } + + @Test + public void testActiveParentControllerPasses() { + when(admin.isParent()).thenReturn(true); + when(admin.getParentControllerRegionState()).thenReturn(ParentControllerRegionState.ACTIVE); + + interceptor.interceptCall(call, headers, next); + + verify(next, times(1)).startCall(call, headers); + verify(call, never()).close(any(), any()); + } + + @Test + public void testInactiveParentControllerRejectsRequest() { + when(admin.isParent()).thenReturn(true); + when(admin.getParentControllerRegionState()).thenReturn(ParentControllerRegionState.PASSIVE); + + MethodDescriptor methodDescriptor = VeniceControllerGrpcServiceGrpc.getGetLeaderControllerMethod(); + when(call.getMethodDescriptor()).thenReturn(methodDescriptor); + when(call.getAttributes()).thenReturn(Attributes.EMPTY); + + interceptor.interceptCall(call, headers, next); + + verify(call, times(1)).close(any(io.grpc.Status.class), any(Metadata.class)); + verify(next, never()).startCall(call, headers); + } + + @Test + public void testNonParentControllerPasses() { + when(admin.isParent()).thenReturn(false); + + interceptor.interceptCall(call, headers, next); + + verify(next, times(1)).startCall(call, headers); + verify(call, never()).close(any(), any()); + } + + @Test + public void testErrorMessageAndCodeOnRejection() { + when(admin.isParent()).thenReturn(true); + when(admin.getParentControllerRegionState()).thenReturn(ParentControllerRegionState.PASSIVE); + MethodDescriptor methodDescriptor = VeniceControllerGrpcServiceGrpc.getGetLeaderControllerMethod(); + when(call.getMethodDescriptor()).thenReturn(methodDescriptor); + when(call.getAttributes()).thenReturn(Attributes.EMPTY); + + interceptor.interceptCall(call, headers, next); + + ArgumentCaptor statusCaptor = ArgumentCaptor.forClass(io.grpc.Status.class); + ArgumentCaptor metadataCaptor = ArgumentCaptor.forClass(Metadata.class); + verify(call, times(1)).close(statusCaptor.capture(), metadataCaptor.capture()); + verify(next, never()).startCall(call, headers); + + io.grpc.Status status = statusCaptor.getValue(); + assertTrue(status.getDescription().contains("Parent controller is not active")); + assertEquals(status.getCode(), io.grpc.Status.FAILED_PRECONDITION.getCode()); + } +} diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceDelegateMode.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceDelegateMode.java index 5eea6fdd45a..fa2b6420e4f 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceDelegateMode.java +++ b/services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceDelegateMode.java @@ -512,6 +512,8 @@ public , K, R> Scatter scatter( Map> hostMap = new HashMap<>(); int helixGroupNum = getHelixGroupNum(); int assignedHelixGroupId = getAssignedHelixGroupId(venicePath); + // This is used to record the request start time for the whole Router request. + venicePath.recordOriginalRequestStartTimestamp(); currentPartition = 0; try { for (; currentPartition < partitionCount; currentPartition++) { diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceResponseAggregator.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceResponseAggregator.java index 110469ab30c..f1f77d528ad 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceResponseAggregator.java +++ b/services/venice-router/src/main/java/com/linkedin/venice/router/api/VeniceResponseAggregator.java @@ -181,7 +181,10 @@ public FullHttpResponse buildResponse( * 3. HelixGroupId is valid since Helix-assisted routing is only enabled for multi-key request. */ if (!venicePath.isRetryRequest() && helixGroupSelector != null && venicePath.getHelixGroupId() >= 0) { - helixGroupSelector.finishRequest(venicePath.getRequestId(), venicePath.getHelixGroupId()); + helixGroupSelector.finishRequest( + venicePath.getRequestId(), + venicePath.getHelixGroupId(), + LatencyUtils.getElapsedTimeFromMsToMs(venicePath.getOriginalRequestStartTs())); } RequestType requestType = venicePath.getRequestType(); AggRouterHttpRequestStats stats = routerStats.getStatsByType(requestType); diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupLeastLoadedStrategy.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupLeastLoadedStrategy.java index eaef61b58fb..2d3bdc683fb 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupLeastLoadedStrategy.java +++ b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupLeastLoadedStrategy.java @@ -2,6 +2,7 @@ import com.linkedin.alpini.base.concurrency.TimeoutProcessor; import com.linkedin.venice.exceptions.VeniceException; +import com.linkedin.venice.router.stats.HelixGroupStats; import com.linkedin.venice.utils.Pair; import java.util.HashMap; import java.util.Map; @@ -24,18 +25,18 @@ public class HelixGroupLeastLoadedStrategy implements HelixGroupSelectionStrateg public static final int MAX_ALLOWED_GROUP = 100; private final int[] counters = new int[MAX_ALLOWED_GROUP]; - /** - * The group count could potentially change during the runtime since the storage node cluster can be expanded - * without bouncing Routers. - */ - private int currentGroupCount = 0; private final TimeoutProcessor timeoutProcessor; private final long timeoutInMS; private final Map> requestTimeoutFutureMap = new HashMap<>(); + private final HelixGroupStats helixGroupStats; - public HelixGroupLeastLoadedStrategy(TimeoutProcessor timeoutProcessor, long timeoutInMS) { + public HelixGroupLeastLoadedStrategy( + TimeoutProcessor timeoutProcessor, + long timeoutInMS, + HelixGroupStats helixGroupStats) { this.timeoutProcessor = timeoutProcessor; this.timeoutInMS = timeoutInMS; + this.helixGroupStats = helixGroupStats; } @Override @@ -44,8 +45,8 @@ public int selectGroup(long requestId, int groupCount) { throw new VeniceException( "The valid group num must fail into this range: [1, " + MAX_ALLOWED_GROUP + "], but received: " + groupCount); } - this.currentGroupCount = groupCount; - long smallestCounter = Integer.MAX_VALUE; + int smallestCounter = Integer.MAX_VALUE; + double lowestAvgLatency = Double.MAX_VALUE; int leastLoadedGroup = 0; int startGroupId = (int) (requestId % groupCount); /** @@ -60,10 +61,21 @@ public int selectGroup(long requestId, int groupCount) { } for (int i = 0; i < groupCount; ++i) { int currentGroup = (i + startGroupId) % groupCount; - long currentGroupCounter = counters[currentGroup]; + int currentGroupCounter = counters[currentGroup]; if (currentGroupCounter < smallestCounter) { smallestCounter = currentGroupCounter; leastLoadedGroup = currentGroup; + lowestAvgLatency = helixGroupStats.getGroupResponseWaitingTimeAvg(currentGroup); + } else if (currentGroupCounter == smallestCounter) { + double currentGroupAvgLatency = helixGroupStats.getGroupResponseWaitingTimeAvg(currentGroup); + /** + * Here we don't check whether {@link #currentGroupAvgLatency} is less than 0 or not, as when the group is not + * being used at all, the average latency will be -1.0. + */ + if (currentGroupAvgLatency < lowestAvgLatency) { + lowestAvgLatency = currentGroupAvgLatency; + leastLoadedGroup = currentGroup; + } } } final int finalLeastLoadedGroup = leastLoadedGroup; @@ -82,6 +94,7 @@ public int selectGroup(long requestId, int groupCount) { ++counters[leastLoadedGroup]; } + helixGroupStats.recordGroupPendingRequest(leastLoadedGroup, counters[leastLoadedGroup]); return leastLoadedGroup; } @@ -100,6 +113,10 @@ private void timeoutRequest(long requestId, int groupId, boolean cancelTimeoutFu "The allowed group id must fail into this range: [0, " + (MAX_ALLOWED_GROUP - 1) + "], but received: " + groupId); } + if (!cancelTimeoutFuture) { + // Timeout request + helixGroupStats.recordGroupResponseWaitingTime(groupId, timeoutInMS); + } synchronized (this) { Pair timeoutFuturePair = requestTimeoutFutureMap.get(requestId); if (timeoutFuturePair == null) { @@ -133,47 +150,8 @@ private void timeoutRequest(long requestId, int groupId, boolean cancelTimeoutFu } @Override - public void finishRequest(long requestId, int groupId) { + public void finishRequest(long requestId, int groupId, double latency) { timeoutRequest(requestId, groupId, true); - } - - @Override - public int getMaxGroupPendingRequest() { - if (currentGroupCount == 0) { - return 0; - } - int maxPendingRequest = 0; - for (int i = 0; i < currentGroupCount; ++i) { - if (counters[i] > maxPendingRequest) { - maxPendingRequest = counters[i]; - } - } - return maxPendingRequest; - } - - @Override - public int getMinGroupPendingRequest() { - if (currentGroupCount == 0) { - return 0; - } - int minPendingRequest = Integer.MAX_VALUE; - for (int i = 0; i < currentGroupCount; ++i) { - if (counters[i] < minPendingRequest) { - minPendingRequest = counters[i]; - } - } - return minPendingRequest; - } - - @Override - public int getAvgGroupPendingRequest() { - if (currentGroupCount == 0) { - return 0; - } - int totalPendingRequest = 0; - for (int i = 0; i < currentGroupCount; ++i) { - totalPendingRequest += counters[i]; - } - return totalPendingRequest / currentGroupCount; + helixGroupStats.recordGroupResponseWaitingTime(groupId, latency); } } diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupRoundRobinStrategy.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupRoundRobinStrategy.java index cba210ec7d8..59cb7547ea1 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupRoundRobinStrategy.java +++ b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupRoundRobinStrategy.java @@ -14,25 +14,8 @@ public int selectGroup(long requestId, int groupNum) { } @Override - public void finishRequest(long requestId, int groupId) { + public void finishRequest(long requestId, int groupId, double latency) { // do nothing } - @Override - public int getMaxGroupPendingRequest() { - // Not supported - return -1; - } - - @Override - public int getMinGroupPendingRequest() { - // Not supported - return -1; - } - - @Override - public int getAvgGroupPendingRequest() { - // Not supported - return -1; - } } diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelectionStrategy.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelectionStrategy.java index c60b70b1863..c513e9c604f 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelectionStrategy.java +++ b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelectionStrategy.java @@ -10,20 +10,6 @@ public interface HelixGroupSelectionStrategy { * Notify the corresponding Helix Group that the request is completed, and the implementation will decide whether * any cleanup is required or not. */ - void finishRequest(long requestId, int groupId); + void finishRequest(long requestId, int groupId, double latency); - /** - * Get the maximum of the pending requests among all the groups - */ - int getMaxGroupPendingRequest(); - - /** - * Get the minimum of the pending requests among all the groups - */ - int getMinGroupPendingRequest(); - - /** - * Get the average of the pending requests among all the groups - */ - int getAvgGroupPendingRequest(); } diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelector.java b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelector.java index c36023b87b1..36eb2529e13 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelector.java +++ b/services/venice-router/src/main/java/com/linkedin/venice/router/api/routing/helix/HelixGroupSelector.java @@ -30,11 +30,12 @@ public HelixGroupSelector( HelixInstanceConfigRepository instanceConfigRepository, HelixGroupSelectionStrategyEnum strategyEnum, TimeoutProcessor timeoutProcessor) { - this.helixGroupStats = new HelixGroupStats(metricsRepository, this); + this.helixGroupStats = new HelixGroupStats(metricsRepository); this.instanceConfigRepository = instanceConfigRepository; Class strategyClass = strategyEnum.getStrategyClass(); if (strategyClass.equals(HelixGroupLeastLoadedStrategy.class)) { - this.selectionStrategy = new HelixGroupLeastLoadedStrategy(timeoutProcessor, HELIX_GROUP_COUNTER_TIMEOUT_MS); + this.selectionStrategy = + new HelixGroupLeastLoadedStrategy(timeoutProcessor, HELIX_GROUP_COUNTER_TIMEOUT_MS, helixGroupStats); } else { try { this.selectionStrategy = strategyClass.getDeclaredConstructor().newInstance(); @@ -57,27 +58,11 @@ public int selectGroup(long requestId, int groupNum) { helixGroupStats.recordGroupNum(groupNum); int assignedGroupId = selectionStrategy.selectGroup(requestId, groupNum); helixGroupStats.recordGroupRequest(assignedGroupId); - return assignedGroupId; } @Override - public void finishRequest(long requestId, int groupId) { - selectionStrategy.finishRequest(requestId, groupId); - } - - @Override - public int getMaxGroupPendingRequest() { - return selectionStrategy.getMaxGroupPendingRequest(); - } - - @Override - public int getMinGroupPendingRequest() { - return selectionStrategy.getMinGroupPendingRequest(); - } - - @Override - public int getAvgGroupPendingRequest() { - return selectionStrategy.getAvgGroupPendingRequest(); + public void finishRequest(long requestId, int groupId, double latency) { + selectionStrategy.finishRequest(requestId, groupId, latency); } } diff --git a/services/venice-router/src/main/java/com/linkedin/venice/router/stats/HelixGroupStats.java b/services/venice-router/src/main/java/com/linkedin/venice/router/stats/HelixGroupStats.java index a248de31cf2..e53d8c57b4e 100644 --- a/services/venice-router/src/main/java/com/linkedin/venice/router/stats/HelixGroupStats.java +++ b/services/venice-router/src/main/java/com/linkedin/venice/router/stats/HelixGroupStats.java @@ -1,35 +1,28 @@ package com.linkedin.venice.router.stats; -import com.linkedin.venice.router.api.routing.helix.HelixGroupSelectionStrategy; import com.linkedin.venice.stats.AbstractVeniceStats; import com.linkedin.venice.utils.concurrent.VeniceConcurrentHashMap; +import io.tehuti.Metric; +import io.tehuti.metrics.MeasurableStat; import io.tehuti.metrics.MetricsRepository; import io.tehuti.metrics.Sensor; -import io.tehuti.metrics.stats.AsyncGauge; import io.tehuti.metrics.stats.Avg; import io.tehuti.metrics.stats.OccurrenceRate; public class HelixGroupStats extends AbstractVeniceStats { private final VeniceConcurrentHashMap groupCounterSensorMap = new VeniceConcurrentHashMap<>(); - private final HelixGroupSelectionStrategy strategy; - + private final VeniceConcurrentHashMap pendingRequestSensorMap = new VeniceConcurrentHashMap<>(); + private final VeniceConcurrentHashMap groupResponseWaitingTimeSensorMap = + new VeniceConcurrentHashMap<>(); + private final VeniceConcurrentHashMap groupResponseWaitingTimeAvgMap = + new VeniceConcurrentHashMap<>(); private final Sensor groupCountSensor; - private final Sensor maxGroupPendingRequest; - private final Sensor minGroupPendingRequest; - private final Sensor avgGroupPendingRequest; - public HelixGroupStats(MetricsRepository metricsRepository, HelixGroupSelectionStrategy strategy) { + public HelixGroupStats(MetricsRepository metricsRepository) { super(metricsRepository, "HelixGroupStats"); - this.strategy = strategy; this.groupCountSensor = registerSensor("group_count", new Avg()); - this.maxGroupPendingRequest = registerSensor( - new AsyncGauge((ignored, ignored2) -> strategy.getMaxGroupPendingRequest(), "max_group_pending_request")); - this.minGroupPendingRequest = registerSensor( - new AsyncGauge((ignored, ignored2) -> strategy.getMinGroupPendingRequest(), "min_group_pending_request")); - this.avgGroupPendingRequest = registerSensor( - new AsyncGauge((ignored, ignored2) -> strategy.getAvgGroupPendingRequest(), "avg_group_pending_request")); } public void recordGroupNum(int groupNum) { @@ -41,4 +34,33 @@ public void recordGroupRequest(int groupId) { .computeIfAbsent(groupId, id -> registerSensor("group_" + groupId + "_request", new OccurrenceRate())); groupSensor.record(); } + + public void recordGroupPendingRequest(int groupId, int pendingRequest) { + Sensor pendingRequestSensor = pendingRequestSensorMap + .computeIfAbsent(groupId, id -> registerSensor("group_" + groupId + "_pending_request", new Avg())); + pendingRequestSensor.record(pendingRequest); + } + + public void recordGroupResponseWaitingTime(int groupId, double responseWaitingTime) { + Sensor groupResponseWaitingTimeSensor = groupResponseWaitingTimeSensorMap.computeIfAbsent(groupId, id -> { + MeasurableStat avg = new Avg(); + Sensor sensor = registerSensor("group_" + groupId + "_response_waiting_time", avg); + groupResponseWaitingTimeAvgMap.put(groupId, getMetricsRepository().getMetric(getMetricFullName(sensor, avg))); + + return sensor; + }); + groupResponseWaitingTimeSensor.record(responseWaitingTime); + } + + public double getGroupResponseWaitingTimeAvg(int groupId) { + Metric groupResponseWaitingTimeAvgMetric = groupResponseWaitingTimeAvgMap.get(groupId); + if (groupResponseWaitingTimeAvgMetric == null) { + return -1; + } + double avgLatency = groupResponseWaitingTimeAvgMetric.value(); + if (Double.isNaN(avgLatency)) { + return -1; + } + return avgLatency; + } } diff --git a/services/venice-router/src/test/java/com/linkedin/venice/router/api/path/TestVeniceComputePath.java b/services/venice-router/src/test/java/com/linkedin/venice/router/api/path/TestVeniceComputePath.java index 494f3031421..f4672683dc1 100644 --- a/services/venice-router/src/test/java/com/linkedin/venice/router/api/path/TestVeniceComputePath.java +++ b/services/venice-router/src/test/java/com/linkedin/venice/router/api/path/TestVeniceComputePath.java @@ -1,7 +1,7 @@ package com.linkedin.venice.router.api.path; import static com.linkedin.venice.compute.ComputeRequestWrapper.LATEST_SCHEMA_VERSION_FOR_COMPUTE_REQUEST; -import static org.mockito.ArgumentMatchers.*; +import static org.mockito.ArgumentMatchers.anyString; import static org.mockito.Mockito.any; import static org.mockito.Mockito.anyInt; import static org.mockito.Mockito.mock; diff --git a/services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupLeastLoadedStrategy.java b/services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupLeastLoadedStrategy.java index 9dc92bbe9fd..6a27e7c8112 100644 --- a/services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupLeastLoadedStrategy.java +++ b/services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupLeastLoadedStrategy.java @@ -6,6 +6,8 @@ import static org.mockito.Mockito.mock; import com.linkedin.alpini.base.concurrency.TimeoutProcessor; +import com.linkedin.venice.router.stats.HelixGroupStats; +import io.tehuti.metrics.MetricsRepository; import org.testng.Assert; import org.testng.annotations.Test; @@ -15,24 +17,39 @@ public class TestHelixGroupLeastLoadedStrategy { public void testSelectGroup() { TimeoutProcessor timeoutProcessor = mock(TimeoutProcessor.class); doReturn(mock(TimeoutProcessor.TimeoutFuture.class)).when(timeoutProcessor).schedule(any(), anyLong(), any()); - HelixGroupLeastLoadedStrategy strategy = new HelixGroupLeastLoadedStrategy(timeoutProcessor, 10000); + HelixGroupLeastLoadedStrategy strategy = + new HelixGroupLeastLoadedStrategy(timeoutProcessor, 10000, mock(HelixGroupStats.class)); int groupNum = 3; // Group 0 is slow. Assert.assertEquals(strategy.selectGroup(0, groupNum), 0); Assert.assertEquals(strategy.selectGroup(1, groupNum), 1); Assert.assertEquals(strategy.selectGroup(2, groupNum), 2); - Assert.assertEquals(strategy.getMaxGroupPendingRequest(), 1); - Assert.assertEquals(strategy.getMinGroupPendingRequest(), 1); - Assert.assertEquals(strategy.getAvgGroupPendingRequest(), 1); - strategy.finishRequest(1, 1); - strategy.finishRequest(2, 2); + strategy.finishRequest(1, 1, 1); + strategy.finishRequest(2, 2, 1); Assert.assertEquals(strategy.selectGroup(3, groupNum), 1); Assert.assertEquals(strategy.selectGroup(4, groupNum), 2); - strategy.finishRequest(0, 0); - strategy.finishRequest(3, 1); - strategy.finishRequest(4, 2); + strategy.finishRequest(0, 0, 1); + strategy.finishRequest(3, 1, 1); + strategy.finishRequest(4, 2, 1); // Group 0 is recovered Assert.assertEquals(strategy.selectGroup(5, groupNum), 2); Assert.assertEquals(strategy.selectGroup(6, groupNum), 0); } + + @Test + public void testLatencyBasedGroupSelection() { + TimeoutProcessor timeoutProcessor = mock(TimeoutProcessor.class); + doReturn(mock(TimeoutProcessor.TimeoutFuture.class)).when(timeoutProcessor).schedule(any(), anyLong(), any()); + HelixGroupStats stats = new HelixGroupStats(new MetricsRepository()); + HelixGroupLeastLoadedStrategy strategy = new HelixGroupLeastLoadedStrategy(timeoutProcessor, 10000, stats); + int groupNum = 3; + Assert.assertEquals(strategy.selectGroup(0, groupNum), 0); + Assert.assertEquals(strategy.selectGroup(1, groupNum), 1); + Assert.assertEquals(strategy.selectGroup(2, groupNum), 2); + // Group 2 is the fastest one + strategy.finishRequest(0, 0, 2); + strategy.finishRequest(1, 1, 3); + strategy.finishRequest(2, 2, 1); + Assert.assertEquals(strategy.selectGroup(3, groupNum), 2); + } } diff --git a/services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupRoundRobinStrategy.java b/services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupRoundRobinStrategy.java index 9d2a877d0c4..c98c91759a3 100644 --- a/services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupRoundRobinStrategy.java +++ b/services/venice-router/src/test/java/com/linkedin/venice/router/api/routing/helix/TestHelixGroupRoundRobinStrategy.java @@ -10,16 +10,16 @@ public void testSelectGroup() { HelixGroupRoundRobinStrategy strategy = new HelixGroupRoundRobinStrategy(); int groupNum = 3; Assert.assertEquals(strategy.selectGroup(0, groupNum), 0); - strategy.finishRequest(0, 0); + strategy.finishRequest(0, 0, 1); Assert.assertEquals(strategy.selectGroup(1, groupNum), 1); - strategy.finishRequest(1, 1); + strategy.finishRequest(1, 1, 1); Assert.assertEquals(strategy.selectGroup(2, groupNum), 2); - strategy.finishRequest(2, 2); + strategy.finishRequest(2, 2, 1); Assert.assertEquals(strategy.selectGroup(3, groupNum), 0); - strategy.finishRequest(3, 0); + strategy.finishRequest(3, 0, 1); Assert.assertEquals(strategy.selectGroup(4, groupNum), 1); - strategy.finishRequest(4, 1); + strategy.finishRequest(4, 1, 1); Assert.assertEquals(strategy.selectGroup(5, groupNum), 2); - strategy.finishRequest(5, 2); + strategy.finishRequest(5, 2, 1); } } diff --git a/services/venice-router/src/test/java/com/linkedin/venice/router/stats/AdminOperationsStatsTest.java b/services/venice-router/src/test/java/com/linkedin/venice/router/stats/AdminOperationsStatsTest.java index f420e0bf253..4b4fee40fcd 100644 --- a/services/venice-router/src/test/java/com/linkedin/venice/router/stats/AdminOperationsStatsTest.java +++ b/services/venice-router/src/test/java/com/linkedin/venice/router/stats/AdminOperationsStatsTest.java @@ -1,6 +1,7 @@ package com.linkedin.venice.router.stats; -import static org.mockito.Mockito.*; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; import com.linkedin.venice.router.VeniceRouterConfig; import com.linkedin.venice.tehuti.MockTehutiReporter; diff --git a/services/venice-server/src/main/java/com/linkedin/venice/listener/ServerReadMetadataRepository.java b/services/venice-server/src/main/java/com/linkedin/venice/listener/ServerReadMetadataRepository.java index 878e231566e..39fdaacce0d 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/listener/ServerReadMetadataRepository.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/listener/ServerReadMetadataRepository.java @@ -180,7 +180,7 @@ public StorePropertiesResponse getStoreProperties(String storeName, Optional entry: valueSchemas.entrySet()) { int schemaId = Integer.parseInt(entry.getKey().toString()); if (schemaId > largestKnownSchemaId.get()) { - storeValueSchemas.put(schemaId, entry.getValue()); + storeValueSchemas.valueSchemaMap.put(Integer.toString(schemaId), entry.getValue()); } } } else { diff --git a/services/venice-server/src/main/java/com/linkedin/venice/server/VeniceServer.java b/services/venice-server/src/main/java/com/linkedin/venice/server/VeniceServer.java index f9d625c8379..ae431565821 100644 --- a/services/venice-server/src/main/java/com/linkedin/venice/server/VeniceServer.java +++ b/services/venice-server/src/main/java/com/linkedin/venice/server/VeniceServer.java @@ -3,11 +3,13 @@ import com.linkedin.avro.fastserde.FastDeserializerGeneratorAccessor; import com.linkedin.davinci.blobtransfer.BlobTransferManager; import com.linkedin.davinci.blobtransfer.BlobTransferUtil; +import com.linkedin.davinci.blobtransfer.BlobTransferUtils.BlobTransferTableFormat; import com.linkedin.davinci.compression.StorageEngineBackedCompressorFactory; import com.linkedin.davinci.config.VeniceClusterConfig; import com.linkedin.davinci.config.VeniceConfigLoader; import com.linkedin.davinci.config.VeniceServerConfig; import com.linkedin.davinci.helix.HelixParticipationService; +import com.linkedin.davinci.kafka.consumer.AdaptiveThrottlerSignalService; import com.linkedin.davinci.kafka.consumer.KafkaStoreIngestionService; import com.linkedin.davinci.kafka.consumer.RemoteIngestionRepairService; import com.linkedin.davinci.repository.VeniceMetadataRepositoryBuilder; @@ -118,6 +120,7 @@ public class VeniceServer { private ICProvider icProvider; StorageEngineBackedCompressorFactory compressorFactory; private HeartbeatMonitoringService heartbeatMonitoringService; + private AdaptiveThrottlerSignalService adaptiveThrottlerSignalService; private ServerReadMetadataRepository serverReadMetadataRepository; private BlobTransferManager blobTransferManager; private AggVersionedBlobTransferStats aggVersionedBlobTransferStats; @@ -374,6 +377,12 @@ private List createServices() { services.add(heartbeatMonitoringService); this.zkHelixAdmin = Lazy.of(() -> new ZKHelixAdmin(serverConfig.getZookeeperAddress())); + this.adaptiveThrottlerSignalService = null; + if (serverConfig.isAdaptiveThrottlerEnabled()) { + adaptiveThrottlerSignalService = + new AdaptiveThrottlerSignalService(serverConfig, metricsRepository, heartbeatMonitoringService); + services.add(adaptiveThrottlerSignalService); + } // create and add KafkaSimpleConsumerService this.kafkaStoreIngestionService = new KafkaStoreIngestionService( storageService, @@ -398,7 +407,8 @@ private List createServices() { pubSubClientsFactory, sslFactory, heartbeatMonitoringService, - zkHelixAdmin); + zkHelixAdmin, + adaptiveThrottlerSignalService); this.diskHealthCheckService = new DiskHealthCheckService( serverConfig.isDiskHealthCheckServiceEnabled(), @@ -473,7 +483,10 @@ private List createServices() { serverConfig.getMaxConcurrentSnapshotUser(), serverConfig.getSnapshotRetentionTimeInMin(), serverConfig.getBlobTransferMaxTimeoutInMin(), - aggVersionedBlobTransferStats); + aggVersionedBlobTransferStats, + serverConfig.getRocksDBServerConfig().isRocksDBPlainTableFormatEnabled() + ? BlobTransferTableFormat.PLAIN_TABLE + : BlobTransferTableFormat.BLOCK_BASED_TABLE); } else { aggVersionedBlobTransferStats = null; blobTransferManager = null; diff --git a/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerConfigTest.java b/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerConfigTest.java deleted file mode 100644 index e99d681cd84..00000000000 --- a/services/venice-server/src/test/java/com/linkedin/venice/grpc/VeniceGrpcServerConfigTest.java +++ /dev/null @@ -1,86 +0,0 @@ -package com.linkedin.venice.grpc; - -import static org.mockito.Mockito.mock; -import static org.testng.Assert.assertEquals; -import static org.testng.Assert.assertNotNull; -import static org.testng.Assert.assertNull; - -import com.linkedin.alpini.base.concurrency.ExecutorService; -import com.linkedin.venice.security.SSLFactory; -import io.grpc.BindableService; -import io.grpc.ServerCredentials; -import io.grpc.ServerInterceptor; -import java.util.concurrent.Executor; -import org.testng.annotations.Test; - - -public class VeniceGrpcServerConfigTest { - @Test - public void testDefaults() { - VeniceGrpcServerConfig config = new VeniceGrpcServerConfig.Builder().setPort(8080) - .setService(mock(BindableService.class)) - .setNumThreads(10) - .build(); - - assertEquals(config.getPort(), 8080); - assertNull(config.getCredentials()); - assertEquals(config.getInterceptors().size(), 0); - } - - @Test - public void testCustomCredentials() { - VeniceGrpcServerConfig config = new VeniceGrpcServerConfig.Builder().setPort(8080) - .setService(mock(BindableService.class)) - .setCredentials(mock(ServerCredentials.class)) - .setExecutor(mock(ExecutorService.class)) - .setNumThreads(10) - .build(); - - assertNotNull(config.getCredentials()); - assertEquals(config.getCredentials(), config.getCredentials()); - } - - @Test - public void testInterceptor() { - ServerInterceptor interceptor = mock(ServerInterceptor.class); - VeniceGrpcServerConfig config = new VeniceGrpcServerConfig.Builder().setPort(8080) - .setService(mock(BindableService.class)) - .setInterceptor(interceptor) - .setNumThreads(10) - .build(); - - assertEquals(config.getInterceptors().size(), 1); - assertEquals(config.getInterceptors().get(0), interceptor); - } - - @Test - public void testSSLFactory() { - SSLFactory sslFactory = mock(SSLFactory.class); - VeniceGrpcServerConfig config = new VeniceGrpcServerConfig.Builder().setPort(8080) - .setService(mock(BindableService.class)) - .setSslFactory(sslFactory) - .setNumThreads(10) - .build(); - - assertEquals(config.getSslFactory(), sslFactory); - } - - @Test - public void testNumThreadsAndExecutor() { - VeniceGrpcServerConfig.Builder configBuilder = - new VeniceGrpcServerConfig.Builder().setPort(1010).setService(mock(BindableService.class)).setNumThreads(10); - - VeniceGrpcServerConfig testExectorCreation = configBuilder.build(); - - Executor exec = testExectorCreation.getExecutor(); - assertNotNull(exec); - - VeniceGrpcServerConfig testCustomExecutor = configBuilder.setExecutor(exec).build(); - assertEquals(testCustomExecutor.getExecutor(), exec); - } - - @Test(expectedExceptions = IllegalArgumentException.class) - public void testNoService() { - new VeniceGrpcServerConfig.Builder().setPort(8080).build(); - } -} diff --git a/settings.gradle b/settings.gradle index 7a52d4739f3..834f01adb53 100644 --- a/settings.gradle +++ b/settings.gradle @@ -6,6 +6,10 @@ pluginManagement { url "https://linkedin.jfrog.io/artifactory/open-source" } gradlePluginPortal() + maven { + // Needed for DuckDB SNAPSHOT. TODO: Remove when the real release is published! + url = uri('https://oss.sonatype.org/content/repositories/snapshots/') + } } } @@ -65,6 +69,7 @@ include 'internal:alpini:router:alpini-router-impl' // 3rd-party system integration modules include 'integrations:venice-beam' +include 'integrations:venice-duckdb' include 'integrations:venice-pulsar' include 'integrations:venice-samza'