This is an automated email from the ASF dual-hosted git repository. roryqi pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit d5743b19eeddecf06e4ea5ad9ce9e84670ccc8ab Author: MaSai <[email protected]> AuthorDate: Mon Oct 26 10:36:54 2026 +0800 [#14205] feat(tms): Before/after metrics and finish REST for in-job actions (#14121) Second slice of the in-job track (A2). Stacked on A1 (#14054); targets `branch-tms-scheduler`. - REST callbacks for in-job Spark: `metrics/before`, `metrics/after`, and `finish` (`TableMaintenanceActionOperations`) - `MaintenanceMetricsService`, binder, and REST utils (aligned with idp-basic style) - Shared request DTOs in new module `:maintenance:maintenance-common` - OpenAPI under `docs/open-api/maintenance/` - One Spark job may bind multiple tables / soft slots; callbacks key by `(jobId, tableId, policyId)` - Lazy-bind when a Spark callback arrives before submit-side `bindJobId` - Resolve metalake from the **table id** via `TableMetaService.getMetalakeNameByTableId` In-job Spark jobs need HTTP hooks to record before/after snapshot metrics and to release soft occupancy per table (with or without metrics). Wire DTOs belong in a thin shared `*-common` module so the TMS plugin and built-in jobs can share them without coupling through `client-java`. Fix: #14205 Part of #13760. Depends on #14054. Yes — maintenance action metrics/finish REST (plugin OpenAPI): - `POST /metalakes/{metalake}/maintenance/actions/{jobId}/metrics/before` - `POST /metalakes/{metalake}/maintenance/actions/{jobId}/metrics/after` - `POST /metalakes/{metalake}/maintenance/actions/{jobId}/finish` `./gradlew :maintenance:maintenance-common:compileJava :plugins:tms:compileJava :plugins:tms:compileTestJava -PskipITs` `./gradlew :plugins:tms:test -PskipITs -PskipDockerTests=true --tests org.apache.gravitino.maintenance.web.rest.TestTableMaintenanceActionOperations` --------- Co-authored-by: Cursor <[email protected]> --- .github/workflows/build.yml | 1 + build.gradle.kts | 1 + design-docs/table-maintenance-service-design.md | 4 +- docs/build.gradle.kts | 12 + docs/open-api/maintenance/maintenance.yaml | 224 ++++++++++++++++ docs/open-api/maintenance/openapi.yaml | 69 +++++ .../maintenance-common}/build.gradle.kts | 47 ++-- .../requests/MaintenanceActionFinishRequest.java | 88 +++++++ .../requests/MaintenanceActionMetricsRequest.java | 82 ++++++ plugins/tms/build.gradle.kts | 1 + .../maintenance/MaintenanceMetricsService.java | 218 ++++++++++++++++ .../injob/TableMaintenanceInJobSampler.java | 2 +- .../maintenance/occupancy/OccupancyGate.java | 169 +----------- .../mapper/TableMaintenanceActionMapper.java | 34 ++- .../TableMaintenanceActionSQLProviderFactory.java | 55 +++- .../TableMaintenanceActionBaseSQLProvider.java | 66 ++++- .../service/TableMaintenanceActionMetaService.java | 170 ++++++++++-- .../maintenance/web/TableMaintenanceRESTUtils.java | 184 +++++++++++++ .../web/rest/TableMaintenanceActionOperations.java | 188 ++++++++++++++ .../web/rest/TableMaintenanceBinder.java | 35 +-- .../rest/feature/TableMaintenanceRESTFeature.java | 22 +- .../maintenance/occupancy/TestOccupancyGate.java | 286 ++------------------- .../TestTableMaintenanceActionMetaService.java | 189 +++++++++++--- .../rest/TestTableMaintenanceActionOperations.java | 234 +++++++++++++++++ scripts/h2/schema-2.0.0-h2.sql | 2 +- scripts/h2/upgrade-1.3.0-to-2.0.0-h2.sql | 2 +- scripts/mysql/schema-2.0.0-mysql.sql | 2 +- scripts/mysql/upgrade-1.3.0-to-2.0.0-mysql.sql | 2 +- scripts/postgresql/schema-2.0.0-postgresql.sql | 2 +- .../upgrade-1.3.0-to-2.0.0-postgresql.sql | 2 +- settings.gradle.kts | 1 + 31 files changed, 1823 insertions(+), 571 deletions(-) diff --git a/.github/workflows/build.yml b/.github/workflows/build.yml index 8414697266..7808018292 100644 --- a/.github/workflows/build.yml +++ b/.github/workflows/build.yml @@ -239,6 +239,7 @@ jobs: run: | if [ "${{ needs.changes.outputs.maintenance_module_only_changes }}" = "true" ]; then ./gradlew \ + :maintenance:maintenance-common:build \ :maintenance:optimizer-api:build \ :maintenance:updaters:build \ :maintenance:optimizer:build \ diff --git a/build.gradle.kts b/build.gradle.kts index a1295f06c3..c1201dfb18 100644 --- a/build.gradle.kts +++ b/build.gradle.kts @@ -625,6 +625,7 @@ subprojects { ":catalogs:hadoop-auth", ":catalogs:hadoop-common", ":maintenance:jobs", + ":maintenance:maintenance-common", ":maintenance:optimizer-api", ":maintenance:updaters", ":clients", diff --git a/design-docs/table-maintenance-service-design.md b/design-docs/table-maintenance-service-design.md index 5ec6be113f..5913c07926 100644 --- a/design-docs/table-maintenance-service-design.md +++ b/design-docs/table-maintenance-service-design.md @@ -289,7 +289,7 @@ gates**, **`minIntervalMs`**, and Iceberg **before/after snapshot ids**. That is | Column | Type | Notes | | -------------------- | -------------------------- | ------------------------------------------------------------------------------------------ | | `id` | `BIGINT UNSIGNED` PK | Auto increment | -| `job_id` | `BIGINT UNSIGNED NULL` | Set at `runJob`; maps to `job_run_meta.job_run_id`; null if released before submit | +| `job_id` | `BIGINT UNSIGNED NULL` | Set at `runJob`; maps to `job_run_meta.job_run_id`; **non-unique** — one Spark job may bind multiple tables; null if released before submit | | `table_id` | `BIGINT UNSIGNED NOT NULL` | `table_meta` PK | | `policy_id` | `BIGINT UNSIGNED NOT NULL` | Cron: that Job's policy. Commit chain: first type that ran; full chain in runtime template | | `before_snapshot_id` | `BIGINT NULL` | Set inside the Spark Job when S0 is sampled | @@ -306,7 +306,7 @@ CREATE TABLE IF NOT EXISTS `table_maintenance_job` ( `after_snapshot_id` BIGINT(20) NULL COMMENT 'Iceberg snapshot id after job; null while pending', `finished_at` BIGINT(20) UNSIGNED NULL COMMENT 'end time; null = in-flight', PRIMARY KEY (`id`), - UNIQUE KEY `uk_job_id` (`job_id`), + KEY `idx_job_id` (`job_id`), KEY `idx_tmj_table_policy_finished` (`table_id`, `policy_id`, `finished_at`), KEY `idx_tmj_before_sid` (`table_id`, `before_snapshot_id`), KEY `idx_tmj_after_sid` (`table_id`, `after_snapshot_id`) diff --git a/docs/build.gradle.kts b/docs/build.gradle.kts index 39810c3dda..302aa58773 100644 --- a/docs/build.gradle.kts +++ b/docs/build.gradle.kts @@ -43,8 +43,20 @@ tasks { ) } + val lintMaintenanceOpenAPI by registering(NpxTask::class) { + command.set("@redocly/[email protected]") + args.set( + listOf( + "lint", + "--extends=recommended-strict", + "${project.projectDir}/open-api/maintenance/openapi.yaml" + ) + ) + } + build { dependsOn(lintOpenAPI) dependsOn(lintIdpOpenAPI) + dependsOn(lintMaintenanceOpenAPI) } } diff --git a/docs/open-api/maintenance/maintenance.yaml b/docs/open-api/maintenance/maintenance.yaml new file mode 100644 index 0000000000..7ac6375731 --- /dev/null +++ b/docs/open-api/maintenance/maintenance.yaml @@ -0,0 +1,224 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +--- + +paths: + + /metalakes/{metalake}/maintenance/actions/{jobId}/metrics/before: + parameters: + - $ref: "../openapi.yaml#/components/parameters/metalake" + - $ref: "#/components/parameters/jobId" + - $ref: "#/components/parameters/finish" + post: + tags: + - maintenance + summary: Report TMS before-maintenance (S0) snapshot metrics + operationId: reportMaintenanceActionBeforeMetrics + description: | + Spark callback for an in-flight maintenance action at S0: persists snapshot metrics and + records `before_snapshot_id`. `tableId` and `policyId` select which bound action row when + one Spark job covers multiple tables and/or policies. Optional query `finish=true` closes + that table/policy soft slot (for example on snapshot drift / skip). Continue / stop stays + in the job. + requestBody: + required: true + content: + application/vnd.gravitino.v1+json: + schema: + $ref: "#/components/schemas/MaintenanceActionMetricsRequest" + examples: + MaintenanceActionMetricsRequest: + $ref: "#/components/examples/MaintenanceActionBeforeMetricsRequest" + responses: + "200": + $ref: "../openapi.yaml#/components/responses/BaseResponse" + "400": + $ref: "../openapi.yaml#/components/responses/BadRequestErrorResponse" + "404": + description: Not Found - The specified metalake or maintenance action does not exist + content: + application/vnd.gravitino.v1+json: + schema: + $ref: "../openapi.yaml#/components/schemas/ErrorModel" + "5xx": + $ref: "../openapi.yaml#/components/responses/ServerErrorResponse" + + /metalakes/{metalake}/maintenance/actions/{jobId}/finish: + parameters: + - $ref: "../openapi.yaml#/components/parameters/metalake" + - $ref: "#/components/parameters/jobId" + post: + tags: + - maintenance + summary: Finish TMS maintenance action soft slot without metrics + operationId: finishMaintenanceAction + description: | + Closes the soft slot for one bound table/policy without writing snapshot metrics. + Use when the maintenance type does not sample metrics, or when metrics were already + reported and only occupancy release remains. `tableId` and `policyId` select the action + row. Optional `afterSnapshotId` sets `after_snapshot_id` when known. + requestBody: + required: true + content: + application/vnd.gravitino.v1+json: + schema: + $ref: "#/components/schemas/MaintenanceActionFinishRequest" + examples: + MaintenanceActionFinishRequest: + $ref: "#/components/examples/MaintenanceActionFinishRequest" + responses: + "200": + $ref: "../openapi.yaml#/components/responses/BaseResponse" + "400": + $ref: "../openapi.yaml#/components/responses/BadRequestErrorResponse" + "404": + description: Not Found - The specified metalake or maintenance action does not exist + content: + application/vnd.gravitino.v1+json: + schema: + $ref: "../openapi.yaml#/components/schemas/ErrorModel" + "5xx": + $ref: "../openapi.yaml#/components/responses/ServerErrorResponse" + + /metalakes/{metalake}/maintenance/actions/{jobId}/metrics/after: + parameters: + - $ref: "../openapi.yaml#/components/parameters/metalake" + - $ref: "#/components/parameters/jobId" + - $ref: "#/components/parameters/finish" + post: + tags: + - maintenance + summary: Report TMS after-maintenance (S1) snapshot metrics + operationId: reportMaintenanceActionAfterMetrics + description: | + Spark callback for an in-flight maintenance action at S1: persists snapshot metrics. + `tableId` and `policyId` select which bound action row when one Spark job covers multiple + tables and/or policies. Optional query `finish=true` closes that table/policy soft slot + and sets `after_snapshot_id` to `snapshotId`. Continue / stop stays in the job. + requestBody: + required: true + content: + application/vnd.gravitino.v1+json: + schema: + $ref: "#/components/schemas/MaintenanceActionMetricsRequest" + examples: + MaintenanceActionMetricsRequest: + $ref: "#/components/examples/MaintenanceActionAfterMetricsRequest" + responses: + "200": + $ref: "../openapi.yaml#/components/responses/BaseResponse" + "400": + $ref: "../openapi.yaml#/components/responses/BadRequestErrorResponse" + "404": + description: Not Found - The specified metalake or maintenance action does not exist + content: + application/vnd.gravitino.v1+json: + schema: + $ref: "../openapi.yaml#/components/schemas/ErrorModel" + "5xx": + $ref: "../openapi.yaml#/components/responses/ServerErrorResponse" + +components: + + parameters: + jobId: + name: jobId + in: path + description: Gravitino job run id (numeric) + required: true + schema: + type: integer + format: int64 + finish: + name: finish + in: query + description: | + When true, close the soft slot for this job/table/policy after storing metrics. + Before: typically on drift/skip (after_snapshot_id left null). After: typically after + successful maintenance (after_snapshot_id = snapshotId). Prefer POST .../finish when no + metrics are written. + required: false + schema: + type: boolean + default: false + + schemas: + MaintenanceActionMetricsRequest: + type: object + required: + - tableId + - policyId + - snapshotId + - metricsJson + properties: + tableId: + type: integer + format: int64 + description: Table entity id bound to this job + policyId: + type: integer + format: int64 + description: Policy entity id bound to this job (required when one job covers multiple policies) + snapshotId: + type: integer + format: int64 + description: Iceberg snapshot id the metrics were collected for + metricsJson: + type: string + description: Serialized metrics JSON for table_snapshot_metrics + MaintenanceActionFinishRequest: + type: object + required: + - tableId + - policyId + properties: + tableId: + type: integer + format: int64 + description: Table entity id bound to this job + policyId: + type: integer + format: int64 + description: Policy entity id bound to this job + afterSnapshotId: + type: integer + format: int64 + nullable: true + description: Optional terminal snapshot id stored as after_snapshot_id + + examples: + MaintenanceActionBeforeMetricsRequest: + value: { + "tableId": 10, + "policyId": 20, + "snapshotId": 9001, + "metricsJson": "{\"maintenanceType\":\"rewrite_data_files\",\"snapshotId\":9001,\"custom-data-file-mse\":1.5}" + } + MaintenanceActionAfterMetricsRequest: + value: { + "tableId": 10, + "policyId": 20, + "snapshotId": 9002, + "metricsJson": "{\"maintenanceType\":\"rewrite_data_files\",\"snapshotId\":9002,\"custom-data-file-mse\":0.2}" + } + MaintenanceActionFinishRequest: + value: { + "tableId": 10, + "policyId": 20, + "afterSnapshotId": null + } diff --git a/docs/open-api/maintenance/openapi.yaml b/docs/open-api/maintenance/openapi.yaml new file mode 100644 index 0000000000..fa137d497e --- /dev/null +++ b/docs/open-api/maintenance/openapi.yaml @@ -0,0 +1,69 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +--- +openapi: 3.0.3 +info: + title: Gravitino Table Maintenance REST API + license: + name: Apache 2.0 + url: https://www.apache.org/licenses/LICENSE-2.0.html + version: 1.3.0-SNAPSHOT + description: | + OpenAPI specification for table-maintenance metrics callbacks exposed by the TMS plugin. + Enable them by registering `org.apache.gravitino.maintenance.web.rest.feature` in + `gravitino.server.rest.extensionPackages`. Callers authenticate with the server + authenticators and must be the metalake owner or hold `RUN_JOB`. + +servers: + - url: "{scheme}://{host}:{port}/{basePath}" + description: Generic base server URL, with all parts configurable. + variables: + scheme: + description: The scheme of the URI, either http or https. + default: http + host: + description: The host address for the specified server + default: localhost + port: + description: The port used when addressing the host + default: "8090" + basePath: + description: Optional prefix to be appended to all routes + default: "api" + +security: + - OAuth2WithJWT: [serviceAudience] + - BasicAuth: [] + - KerberosAuth: [] + +paths: + /metalakes/{metalake}/maintenance/actions/{jobId}/metrics/before: + $ref: "./maintenance.yaml#/paths/~1metalakes~1%7Bmetalake%7D~1maintenance~1actions~1%7BjobId%7D~1metrics~1before" + /metalakes/{metalake}/maintenance/actions/{jobId}/metrics/after: + $ref: "./maintenance.yaml#/paths/~1metalakes~1%7Bmetalake%7D~1maintenance~1actions~1%7BjobId%7D~1metrics~1after" + /metalakes/{metalake}/maintenance/actions/{jobId}/finish: + $ref: "./maintenance.yaml#/paths/~1metalakes~1%7Bmetalake%7D~1maintenance~1actions~1%7BjobId%7D~1finish" + +components: + securitySchemes: + OAuth2WithJWT: + $ref: "../openapi.yaml#/components/securitySchemes/OAuth2WithJWT" + BasicAuth: + $ref: "../openapi.yaml#/components/securitySchemes/BasicAuth" + KerberosAuth: + $ref: "../openapi.yaml#/components/securitySchemes/KerberosAuth" diff --git a/docs/build.gradle.kts b/maintenance/maintenance-common/build.gradle.kts similarity index 53% copy from docs/build.gradle.kts copy to maintenance/maintenance-common/build.gradle.kts index 39810c3dda..9bcab80a90 100644 --- a/docs/build.gradle.kts +++ b/maintenance/maintenance-common/build.gradle.kts @@ -16,35 +16,30 @@ * specific language governing permissions and limitations * under the License. */ +description = "Gravitino Maintenance shared common (TMS wire types)" -import com.github.gradle.node.NodeExtension -import com.github.gradle.node.npm.task.NpxTask - -configure<NodeExtension> { - version = "21.6.1" - npmVersion = "10.2.4" - download = true +plugins { + `maven-publish` + id("java") + id("idea") } -tasks { - val lintOpenAPI by registering(NpxTask::class) { - command.set("@redocly/[email protected]") - args.set(listOf("lint", "--extends=recommended-strict", "${project.projectDir}/open-api/openapi.yaml")) - } +dependencies { + implementation(project(":common")) + + implementation(libs.commons.lang3) + implementation(libs.guava) + implementation(libs.jackson.databind) - val lintIdpOpenAPI by registering(NpxTask::class) { - command.set("@redocly/[email protected]") - args.set( - listOf( - "lint", - "--extends=recommended-strict", - "${project.projectDir}/open-api/idp/openapi.yaml" - ) - ) - } + annotationProcessor(libs.lombok) + compileOnly(libs.lombok) + testAnnotationProcessor(libs.lombok) + testCompileOnly(libs.lombok) + + testImplementation(libs.junit.jupiter.api) + testRuntimeOnly(libs.junit.jupiter.engine) +} - build { - dependsOn(lintOpenAPI) - dependsOn(lintIdpOpenAPI) - } +tasks.test { + useJUnitPlatform() } diff --git a/maintenance/maintenance-common/src/main/java/org/apache/gravitino/maintenance/dto/requests/MaintenanceActionFinishRequest.java b/maintenance/maintenance-common/src/main/java/org/apache/gravitino/maintenance/dto/requests/MaintenanceActionFinishRequest.java new file mode 100644 index 0000000000..414691fbc9 --- /dev/null +++ b/maintenance/maintenance-common/src/main/java/org/apache/gravitino/maintenance/dto/requests/MaintenanceActionFinishRequest.java @@ -0,0 +1,88 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.maintenance.dto.requests; + +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.Preconditions; +import javax.annotation.Nullable; +import lombok.EqualsAndHashCode; +import lombok.Getter; +import org.apache.gravitino.rest.RESTRequest; + +/** + * Closes the soft slot for one table/policy bound to a maintenance job without writing snapshot + * metrics. + * + * <p>Use when a maintenance type does not sample metrics, or when metrics were already reported and + * only occupancy release remains. + */ +@Getter +@EqualsAndHashCode +public class MaintenanceActionFinishRequest implements RESTRequest { + + @JsonProperty("tableId") + private final Long tableId; + + @JsonProperty("policyId") + private final Long policyId; + + /** + * Optional Iceberg snapshot id to store as {@code after_snapshot_id}. Omit or null when releasing + * without an S1 (for example skipped work with no rewrite). + */ + @Nullable + @JsonProperty("afterSnapshotId") + private final Long afterSnapshotId; + + /** Default constructor for Jackson. */ + public MaintenanceActionFinishRequest() { + this(null, null, null); + } + + /** + * Creates a finish request that leaves {@code after_snapshot_id} null. + * + * @param tableId table entity id bound to the job + * @param policyId policy entity id bound to the job + */ + public MaintenanceActionFinishRequest(Long tableId, Long policyId) { + this(tableId, policyId, null); + } + + /** + * Creates a finish request. + * + * @param tableId table entity id bound to the job + * @param policyId policy entity id bound to the job + * @param afterSnapshotId terminal snapshot id, or null + */ + public MaintenanceActionFinishRequest( + Long tableId, Long policyId, @Nullable Long afterSnapshotId) { + this.tableId = tableId; + this.policyId = policyId; + this.afterSnapshotId = afterSnapshotId; + } + + @Override + public void validate() throws IllegalArgumentException { + Preconditions.checkArgument(tableId != null && tableId > 0, "tableId must be a positive long"); + Preconditions.checkArgument( + policyId != null && policyId > 0, "policyId must be a positive long"); + } +} diff --git a/maintenance/maintenance-common/src/main/java/org/apache/gravitino/maintenance/dto/requests/MaintenanceActionMetricsRequest.java b/maintenance/maintenance-common/src/main/java/org/apache/gravitino/maintenance/dto/requests/MaintenanceActionMetricsRequest.java new file mode 100644 index 0000000000..7af5c51997 --- /dev/null +++ b/maintenance/maintenance-common/src/main/java/org/apache/gravitino/maintenance/dto/requests/MaintenanceActionMetricsRequest.java @@ -0,0 +1,82 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.maintenance.dto.requests; + +import com.fasterxml.jackson.annotation.JsonProperty; +import com.google.common.base.Preconditions; +import lombok.EqualsAndHashCode; +import lombok.Getter; +import org.apache.commons.lang3.StringUtils; +import org.apache.gravitino.rest.RESTRequest; + +/** + * Spark callback that persists snapshot metrics for an in-flight maintenance action. + * + * <p>Does not decide whether maintenance should continue — that stays in the Spark job. {@code + * tableId} and {@code policyId} select which bound action row to update when one Spark job covers + * multiple tables and/or policies. Soft-slot release is controlled by the {@code finish} query + * parameter on the metrics endpoint (or by {@code POST .../finish} when no metrics are written). + */ +@Getter +@EqualsAndHashCode +public class MaintenanceActionMetricsRequest implements RESTRequest { + + @JsonProperty("tableId") + private final Long tableId; + + @JsonProperty("policyId") + private final Long policyId; + + @JsonProperty("snapshotId") + private final Long snapshotId; + + @JsonProperty("metricsJson") + private final String metricsJson; + + /** Default constructor for Jackson. */ + public MaintenanceActionMetricsRequest() { + this(null, null, null, null); + } + + /** + * Creates a metrics report. + * + * @param tableId table entity id bound to the job + * @param policyId policy entity id bound to the job + * @param snapshotId Iceberg snapshot id the metrics were collected for + * @param metricsJson serialized metrics for {@code table_snapshot_metrics} + */ + public MaintenanceActionMetricsRequest( + Long tableId, Long policyId, Long snapshotId, String metricsJson) { + this.tableId = tableId; + this.policyId = policyId; + this.snapshotId = snapshotId; + this.metricsJson = metricsJson; + } + + @Override + public void validate() throws IllegalArgumentException { + Preconditions.checkArgument(tableId != null && tableId > 0, "tableId must be a positive long"); + Preconditions.checkArgument( + policyId != null && policyId > 0, "policyId must be a positive long"); + Preconditions.checkArgument(snapshotId != null, "snapshotId must not be null"); + Preconditions.checkArgument( + StringUtils.isNotBlank(metricsJson), "metricsJson must not be blank"); + } +} diff --git a/plugins/tms/build.gradle.kts b/plugins/tms/build.gradle.kts index 33b8a15efc..3f81af0a69 100644 --- a/plugins/tms/build.gradle.kts +++ b/plugins/tms/build.gradle.kts @@ -29,6 +29,7 @@ dependencies { implementation(project(":api")) implementation(project(":common")) implementation(project(":core")) + implementation(project(":maintenance:maintenance-common")) implementation(project(":maintenance:optimizer-api")) implementation(project(":server-common")) diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/MaintenanceMetricsService.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/MaintenanceMetricsService.java new file mode 100644 index 0000000000..94d60a8b1b --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/MaintenanceMetricsService.java @@ -0,0 +1,218 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.maintenance; + +import java.util.function.LongFunction; +import org.apache.gravitino.exceptions.NoSuchEntityException; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionFinishRequest; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionMetricsRequest; +import org.apache.gravitino.maintenance.storage.po.TableMaintenanceActionPO; +import org.apache.gravitino.maintenance.storage.service.TableMaintenanceActionMetaService; +import org.apache.gravitino.maintenance.storage.service.TableSnapshotMetricsMetaService; +import org.apache.gravitino.storage.relational.service.TableMetaService; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Persists before/after snapshot metrics and closes soft slots for in-flight maintenance actions. + * + * <p>Continue / stop stays in the job. Soft-slot release happens via the {@code finish} query + * parameter on metrics reports, or via {@link #finishAction} when no metrics are needed. There is + * no job-terminal reclaim path. + * + * <p>Callbacks that arrive before submit-side {@code bindJobId} lazy-bind the unbound claim for the + * reported {@code (tableId, policyId)} onto the path {@code jobId}. + */ +public class MaintenanceMetricsService { + + private static final Logger LOG = LoggerFactory.getLogger(MaintenanceMetricsService.class); + private static final TableMetaService TABLE_SERVICE = TableMetaService.getInstance(); + + private final TableMaintenanceActionMetaService actionMetaService; + private final TableSnapshotMetricsMetaService metricsMetaService; + private final LongFunction<String> metalakeLookup; + + /** + * Production wiring. + * + * @return service bound into the TMS Jersey feature + */ + public static MaintenanceMetricsService create() { + return new MaintenanceMetricsService( + TableMaintenanceActionMetaService.getInstance(), + TableSnapshotMetricsMetaService.getInstance(), + TABLE_SERVICE::getMetalakeNameByTableId); + } + + /** + * Creates a service with explicit collaborators. + * + * @param actionMetaService action rows + * @param metricsMetaService snapshot metrics rows + * @param metalakeLookup resolves a table id to its metalake name + */ + public MaintenanceMetricsService( + TableMaintenanceActionMetaService actionMetaService, + TableSnapshotMetricsMetaService metricsMetaService, + LongFunction<String> metalakeLookup) { + this.actionMetaService = actionMetaService; + this.metricsMetaService = metricsMetaService; + this.metalakeLookup = metalakeLookup; + } + + /** + * S0 report: stores metrics, records {@code before_snapshot_id}, optionally closes the soft slot + * (for example snapshot drift / skip). + * + * @param metalake metalake name from the request path + * @param jobId Gravitino job run id + * @param finish when {@code true}, close the soft slot after storing metrics + * @param request validated metrics payload + */ + public void reportBeforeMetrics( + String metalake, long jobId, boolean finish, MaintenanceActionMetricsRequest request) { + TableMaintenanceActionPO action = + loadInFlight(metalake, jobId, request.getTableId(), request.getPolicyId()); + metricsMetaService.insertMetrics( + action.tableId(), action.policyId(), request.getSnapshotId(), request.getMetricsJson()); + actionMetaService.recordBeforeSnapshot( + jobId, action.tableId(), action.policyId(), request.getSnapshotId()); + maybeFinish( + jobId, + action, + /*afterSnapshotId=*/ null, + finish, + "before", + metalake, + request.getSnapshotId()); + } + + /** + * S1 report: stores metrics and optionally closes the soft slot with {@code after_snapshot_id} = + * {@code snapshotId}. + * + * @param metalake metalake name from the request path + * @param jobId Gravitino job run id + * @param finish when {@code true}, close the soft slot after storing metrics + * @param request validated metrics payload + */ + public void reportAfterMetrics( + String metalake, long jobId, boolean finish, MaintenanceActionMetricsRequest request) { + TableMaintenanceActionPO action = + loadInFlight(metalake, jobId, request.getTableId(), request.getPolicyId()); + metricsMetaService.insertMetrics( + action.tableId(), action.policyId(), request.getSnapshotId(), request.getMetricsJson()); + maybeFinish( + jobId, action, request.getSnapshotId(), finish, "after", metalake, request.getSnapshotId()); + } + + /** + * Closes the soft slot for one bound table/policy without writing metrics. + * + * @param metalake metalake name from the request path + * @param jobId Gravitino job run id + * @param request finish payload ({@code tableId}, {@code policyId}, optional {@code + * afterSnapshotId}) + */ + public void finishAction(String metalake, long jobId, MaintenanceActionFinishRequest request) { + TableMaintenanceActionPO action = + loadInFlight(metalake, jobId, request.getTableId(), request.getPolicyId()); + actionMetaService.markFinished( + jobId, + action.tableId(), + action.policyId(), + request.getAfterSnapshotId(), + System.currentTimeMillis()); + LOG.info( + "TMS action finished for job {} metalake {} table {} policy {} afterSnapshotId={}", + jobId, + metalake, + action.tableId(), + action.policyId(), + request.getAfterSnapshotId()); + } + + private void maybeFinish( + long jobId, + TableMaintenanceActionPO action, + Long afterSnapshotId, + boolean finish, + String phase, + String metalake, + long snapshotId) { + if (!finish) { + LOG.info( + "TMS {} metrics reported for job {} metalake {} table {} policy {} snapshot {}", + phase, + jobId, + metalake, + action.tableId(), + action.policyId(), + snapshotId); + return; + } + actionMetaService.markFinished( + jobId, action.tableId(), action.policyId(), afterSnapshotId, System.currentTimeMillis()); + LOG.info( + "TMS {} metrics reported and action finished for job {} metalake {} table {} policy {} snapshot {}", + phase, + jobId, + metalake, + action.tableId(), + action.policyId(), + snapshotId); + } + + private TableMaintenanceActionPO loadInFlight( + String metalake, long jobId, long tableId, long policyId) { + TableMaintenanceActionPO action = + actionMetaService.getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyId); + if (action == null) { + if (actionMetaService.bindJobIdIfUnbound(tableId, policyId, jobId)) { + action = actionMetaService.getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyId); + } + } + if (action == null) { + throw new NoSuchEntityException( + "no maintenance action for job_id=%s table_id=%s policy_id=%s", jobId, tableId, policyId); + } + requireActionInFlight(action); + requireActionInMetalake(action, metalake); + return action; + } + + private static void requireActionInFlight(TableMaintenanceActionPO action) { + if (action.finishedAt() == null || action.finishedAt() != 0L) { + throw new NoSuchEntityException( + "no in-flight maintenance action for job_id=%s", action.jobId()); + } + } + + private void requireActionInMetalake(TableMaintenanceActionPO action, String metalake) { + String actionMetalake = metalakeLookup.apply(action.tableId()); + if (actionMetalake == null) { + throw new NoSuchEntityException( + "no metalake for table_id=%s (job_id=%s)", action.tableId(), action.jobId()); + } + if (!actionMetalake.equals(metalake)) { + throw new NoSuchEntityException( + "no maintenance action for job_id=%s under metalake %s", action.jobId(), metalake); + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/TableMaintenanceInJobSampler.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/TableMaintenanceInJobSampler.java index 9c96638434..c6ad0995c1 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/TableMaintenanceInJobSampler.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/TableMaintenanceInJobSampler.java @@ -75,7 +75,7 @@ public class TableMaintenanceInJobSampler { Preconditions.checkNotNull(head, "head"); metricsMetaService.insertMetrics(tableId, policyId, sampledSnapshotId, metricsJson); - actionMetaService.recordBeforeSnapshot(jobId, sampledSnapshotId); + actionMetaService.recordBeforeSnapshot(jobId, tableId, policyId, sampledSnapshotId); long headSnapshotId = head.currentSnapshotId(); if (sampledSnapshotId != headSnapshotId) { diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/occupancy/OccupancyGate.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/occupancy/OccupancyGate.java index cf31ca7502..a817b02f2c 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/occupancy/OccupancyGate.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/occupancy/OccupancyGate.java @@ -19,42 +19,21 @@ package org.apache.gravitino.maintenance.occupancy; import com.google.common.base.Preconditions; -import java.util.EnumSet; -import java.util.Set; import javax.annotation.Nullable; -import org.apache.gravitino.exceptions.NoSuchEntityException; -import org.apache.gravitino.job.JobHandle; import org.apache.gravitino.maintenance.TableMaintenanceConfig; import org.apache.gravitino.maintenance.storage.po.TableMaintenanceActionPO; import org.apache.gravitino.maintenance.storage.service.TableMaintenanceActionMetaService; -import org.apache.gravitino.meta.JobEntity; -import org.apache.gravitino.storage.relational.service.JobMetaService; -import org.apache.gravitino.storage.relational.service.TableMetaService; -import org.apache.gravitino.utils.NameIdentifierUtil; -import org.slf4j.Logger; -import org.slf4j.LoggerFactory; /** * Expand / submit occupancy checks against {@code table_maintenance_action} (§5.2). * * <p>Unbound rows ({@code job_id} null, {@code finished_at = 0}) are a short claim: occupied while - * still within TTL on the <strong>database</strong> clock (same timeline as acquire/reclaim). Bound - * rows consult {@link JobMetaService}; when that job is already terminal or confirmed missing, - * {@link #reclaimTerminalBound} marks {@code finished_at} so the soft slot can be claimed again - * (without inventing stub S1 metrics rows or copying {@code before_snapshot_id} into {@code - * after_snapshot_id}). A later finish listener may still backfill a real S1 via {@link - * TableMaintenanceActionMetaService#markFinished}. + * still within TTL on the <strong>database</strong> clock. Bound rows occupy the soft slot until + * Spark closes them via before/after metrics with {@code ?finish=true}. This gate does not reclaim + * on Gravitino job terminal state. */ public final class OccupancyGate { - private static final Logger LOG = LoggerFactory.getLogger(OccupancyGate.class); - - private static final Set<JobHandle.Status> ACTIVE = - EnumSet.of(JobHandle.Status.QUEUED, JobHandle.Status.STARTED, JobHandle.Status.CANCELLING); - - private static final JobMetaService JOB_SERVICE = JobMetaService.getInstance(); - private static final TableMetaService TABLE_SERVICE = TableMetaService.getInstance(); - private final TableMaintenanceActionMetaService actionMetaService; private final long claimTimeoutMs; @@ -86,14 +65,9 @@ public final class OccupancyGate { * @return true when any open row is a live unbound claim or a still-bound soft-slot row */ public boolean isInFlight(long tableId, long policyId) { - reclaimTerminalBound(tableId, policyId); if (actionMetaService.hasLiveUnboundClaim(tableId, policyId, claimTimeoutMs)) { return true; } - // After reclaim, any remaining bound open row occupies the soft slot (ACTIVE job, unresolved - // lookup, or a terminal/missing job whose markFinished raced/failed). Fail closed so - // UNIQUE(table, policy, finished_at=0) cannot be ignored while still blocking - // tryAcquireInFlight. for (TableMaintenanceActionPO open : actionMetaService.listInFlight(tableId, policyId)) { if (open.jobId() != null) { return true; @@ -102,57 +76,6 @@ public final class OccupancyGate { return false; } - /** - * Job-terminal reconciler failsafe (§5.2.2): when Spark finish callback was missed, mark the - * action finished so the soft slot can be reused. Does not invent stub metrics rows (those would - * pollute {@code table_snapshot_metrics}), and does not copy {@code before_snapshot_id} into - * {@code after_snapshot_id} — only a real S1 (or null) is stored (§5.2.3). - * - * <p>Only reclaim when the bound job is confirmed terminal or confirmed missing. An unresolved - * job lookup (for example the table/metalake cannot be resolved) leaves the row open. - * - * @param tableId table entity id - * @param policyId policy entity id - */ - public void reclaimTerminalBound(long tableId, long policyId) { - for (TableMaintenanceActionPO open : actionMetaService.listInFlight(tableId, policyId)) { - if (open.jobId() == null) { - continue; - } - BoundJobLookup lookup = lookupBoundJob(open.jobId(), tableId); - if (lookup.kind == BoundJobLookup.Kind.ACTIVE - || lookup.kind == BoundJobLookup.Kind.UNRESOLVED) { - if (lookup.kind == BoundJobLookup.Kind.UNRESOLVED) { - LOG.warn( - "Skipping reclaim of job {} for table {} policy {}: metalake could not be resolved", - open.jobId(), - tableId, - policyId); - } - continue; - } - JobEntity job = lookup.job; - long finishedAt = - job != null && job.finishedAt() != null && job.finishedAt() > 0 - ? job.finishedAt() - : System.currentTimeMillis(); - // Keep a real after snapshot only; never invent S1 from S0. - Long afterSnapshotId = open.afterSnapshotId(); - try { - actionMetaService.markFinished(open.jobId(), afterSnapshotId, finishedAt); - } catch (RuntimeException e) { - // Another node may have finished the row, or DB is unhealthy — log and keep fail-closed - // via isInFlight seeing the still-open bound row. - LOG.warn( - "Failed to reclaim terminal/missing job {} for table {} policy {}: {}", - open.jobId(), - tableId, - policyId, - e.toString()); - } - } - } - /** * @param tableId table entity id * @param policyId policy entity id @@ -182,92 +105,8 @@ public final class OccupancyGate { || withinMinInterval(tableId, policyId, minIntervalMs, nowMs); } - /** - * Prefer {@code table_maintenance_action.finished_at}; if that column is still {@code 0} but a - * Gravitino job is already terminal, use {@link JobEntity#finishedAt()} without writing back. - */ @Nullable private Long lastFinishedAt(long tableId, long policyId) { - Long latest = actionMetaService.maxFinishedAt(tableId, policyId); - for (TableMaintenanceActionPO open : actionMetaService.listInFlight(tableId, policyId)) { - if (open.jobId() == null) { - continue; - } - BoundJobLookup lookup = lookupBoundJob(open.jobId(), tableId); - if (lookup.kind != BoundJobLookup.Kind.TERMINAL || lookup.job == null) { - continue; - } - Long fromJob = - lookup.job.finishedAt() != null && lookup.job.finishedAt() > 0 - ? lookup.job.finishedAt() - : null; - if (fromJob == null) { - continue; - } - if (latest == null || fromJob > latest) { - latest = fromJob; - } - } - return latest; - } - - /** - * Resolves the bound job through the table's metalake, not the current policy. A deleted policy - * must not make a still-running job look "missing". - */ - private static BoundJobLookup lookupBoundJob(long jobId, long tableId) { - String metalake = TABLE_SERVICE.getMetalakeNameByTableId(tableId); - if (metalake == null) { - return BoundJobLookup.unresolved(); - } - try { - JobEntity job = - JOB_SERVICE.getJobByIdentifier( - NameIdentifierUtil.ofJob(metalake, JobHandle.JOB_ID_PREFIX + jobId)); - if (ACTIVE.contains(job.status())) { - return BoundJobLookup.active(job); - } - return BoundJobLookup.terminal(job); - } catch (NoSuchEntityException e) { - return BoundJobLookup.missing(); - } - } - - /** Outcome of resolving a bound Gravitino job for reclaim / cooldown. */ - private static final class BoundJobLookup { - enum Kind { - /** Job meta loaded and still queued/started/cancelling. */ - ACTIVE, - /** Job meta loaded and already terminal. */ - TERMINAL, - /** Metalake resolved but job meta is gone. */ - MISSING, - /** Table/metalake could not be resolved; do not treat as terminal. */ - UNRESOLVED - } - - final Kind kind; - @Nullable final JobEntity job; - - private BoundJobLookup(Kind kind, @Nullable JobEntity job) { - this.kind = kind; - this.job = job; - } - - static BoundJobLookup active(JobEntity job) { - return new BoundJobLookup(Kind.ACTIVE, job); - } - - static BoundJobLookup terminal(JobEntity job) { - return new BoundJobLookup(Kind.TERMINAL, job); - } - - static BoundJobLookup missing() { - return new BoundJobLookup(Kind.MISSING, null); - } - - static BoundJobLookup unresolved() { - return new BoundJobLookup(Kind.UNRESOLVED, null); - } + return actionMetaService.maxFinishedAt(tableId, policyId); } } diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionMapper.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionMapper.java index f0acccc30a..c7313611ee 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionMapper.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionMapper.java @@ -52,6 +52,11 @@ public interface TableMaintenanceActionMapper { TableMaintenanceActionPO selectUnboundByTableAndPolicy( @Param("tableId") long tableId, @Param("policyId") long policyId); + @SelectProvider( + type = TableMaintenanceActionSQLProviderFactory.class, + method = "selectUnboundByTableId") + List<TableMaintenanceActionPO> selectUnboundByTableId(@Param("tableId") long tableId); + @Nullable @SelectProvider( type = TableMaintenanceActionSQLProviderFactory.class, @@ -71,19 +76,37 @@ public interface TableMaintenanceActionMapper { @SelectProvider(type = TableMaintenanceActionSQLProviderFactory.class, method = "selectById") TableMaintenanceActionPO selectById(@Param("id") long id); - @Nullable @SelectProvider(type = TableMaintenanceActionSQLProviderFactory.class, method = "selectByJobId") - TableMaintenanceActionPO selectByJobId(@Param("jobId") long jobId); + List<TableMaintenanceActionPO> selectByJobId(@Param("jobId") long jobId); + + @Nullable + @SelectProvider( + type = TableMaintenanceActionSQLProviderFactory.class, + method = "selectByJobIdAndTableId") + TableMaintenanceActionPO selectByJobIdAndTableId( + @Param("jobId") long jobId, @Param("tableId") long tableId); + + @Nullable + @SelectProvider( + type = TableMaintenanceActionSQLProviderFactory.class, + method = "selectByJobIdAndTableIdAndPolicyId") + TableMaintenanceActionPO selectByJobIdAndTableIdAndPolicyId( + @Param("jobId") long jobId, @Param("tableId") long tableId, @Param("policyId") long policyId); @UpdateProvider( type = TableMaintenanceActionSQLProviderFactory.class, method = "updateBeforeSnapshot") int updateBeforeSnapshot( - @Param("jobId") long jobId, @Param("beforeSnapshotId") long beforeSnapshotId); + @Param("jobId") long jobId, + @Param("tableId") long tableId, + @Param("policyId") long policyId, + @Param("beforeSnapshotId") long beforeSnapshotId); @UpdateProvider(type = TableMaintenanceActionSQLProviderFactory.class, method = "markFinished") int markFinished( @Param("jobId") long jobId, + @Param("tableId") long tableId, + @Param("policyId") long policyId, @Param("afterSnapshotId") Long afterSnapshotId, @Param("finishedAt") long finishedAt); @@ -91,7 +114,10 @@ public interface TableMaintenanceActionMapper { type = TableMaintenanceActionSQLProviderFactory.class, method = "backfillAfterSnapshot") int backfillAfterSnapshot( - @Param("jobId") long jobId, @Param("afterSnapshotId") long afterSnapshotId); + @Param("jobId") long jobId, + @Param("tableId") long tableId, + @Param("policyId") long policyId, + @Param("afterSnapshotId") long afterSnapshotId); @DeleteProvider(type = TableMaintenanceActionSQLProviderFactory.class, method = "deleteById") int deleteById(@Param("id") long id); diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionSQLProviderFactory.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionSQLProviderFactory.java index 656d1e2585..784da61e70 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionSQLProviderFactory.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionSQLProviderFactory.java @@ -79,6 +79,14 @@ public class TableMaintenanceActionSQLProviderFactory { return getProvider().selectUnboundByTableAndPolicy(tableId, policyId); } + /** + * @param tableId table entity id + * @return SELECT statement for unbound in-flight rows on the table + */ + public static String selectUnboundByTableId(@Param("tableId") long tableId) { + return getProvider().selectUnboundByTableId(tableId); + } + /** * @param tableId table entity id * @param policyId policy entity id @@ -123,35 +131,72 @@ public class TableMaintenanceActionSQLProviderFactory { /** * @param jobId job run id + * @param tableId table entity id + * @return SELECT statement + */ + public static String selectByJobIdAndTableId( + @Param("jobId") long jobId, @Param("tableId") long tableId) { + return getProvider().selectByJobIdAndTableId(jobId, tableId); + } + + /** + * @param jobId job run id + * @param tableId table entity id + * @param policyId policy entity id + * @return SELECT statement + */ + public static String selectByJobIdAndTableIdAndPolicyId( + @Param("jobId") long jobId, + @Param("tableId") long tableId, + @Param("policyId") long policyId) { + return getProvider().selectByJobIdAndTableIdAndPolicyId(jobId, tableId, policyId); + } + + /** + * @param jobId job run id + * @param tableId table entity id + * @param policyId policy entity id * @param beforeSnapshotId snapshot id * @return UPDATE statement */ public static String updateBeforeSnapshot( - @Param("jobId") long jobId, @Param("beforeSnapshotId") long beforeSnapshotId) { - return getProvider().updateBeforeSnapshot(jobId, beforeSnapshotId); + @Param("jobId") long jobId, + @Param("tableId") long tableId, + @Param("policyId") long policyId, + @Param("beforeSnapshotId") long beforeSnapshotId) { + return getProvider().updateBeforeSnapshot(jobId, tableId, policyId, beforeSnapshotId); } /** * @param jobId job run id + * @param tableId table entity id + * @param policyId policy entity id * @param afterSnapshotId after snapshot * @param finishedAt finish millis * @return UPDATE statement */ public static String markFinished( @Param("jobId") long jobId, + @Param("tableId") long tableId, + @Param("policyId") long policyId, @Param("afterSnapshotId") Long afterSnapshotId, @Param("finishedAt") long finishedAt) { - return getProvider().markFinished(jobId, afterSnapshotId, finishedAt); + return getProvider().markFinished(jobId, tableId, policyId, afterSnapshotId, finishedAt); } /** * @param jobId job run id + * @param tableId table entity id + * @param policyId policy entity id * @param afterSnapshotId after snapshot * @return UPDATE statement */ public static String backfillAfterSnapshot( - @Param("jobId") long jobId, @Param("afterSnapshotId") long afterSnapshotId) { - return getProvider().backfillAfterSnapshot(jobId, afterSnapshotId); + @Param("jobId") long jobId, + @Param("tableId") long tableId, + @Param("policyId") long policyId, + @Param("afterSnapshotId") long afterSnapshotId) { + return getProvider().backfillAfterSnapshot(jobId, tableId, policyId, afterSnapshotId); } /** diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/provider/base/TableMaintenanceActionBaseSQLProvider.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/provider/base/TableMaintenanceActionBaseSQLProvider.java index 457ebddb20..a18f685fbd 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/provider/base/TableMaintenanceActionBaseSQLProvider.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/provider/base/TableMaintenanceActionBaseSQLProvider.java @@ -88,6 +88,17 @@ public class TableMaintenanceActionBaseSQLProvider { + " AND job_id IS NULL AND finished_at = 0"; } + /** + * Unbound in-flight claims for a table (any policy). Used by callback lazy-bind when the Spark + * job reports before submit-side {@code bindJobId}. + * + * @param tableId table entity id + * @return SELECT statement + */ + public String selectUnboundByTableId(@Param("tableId") long tableId) { + return selectColumns() + " WHERE table_id = #{tableId} AND job_id IS NULL AND finished_at = 0"; + } + /** * Live unbound claim still within TTL on the database clock. * @@ -139,47 +150,82 @@ public class TableMaintenanceActionBaseSQLProvider { /** * @param jobId Gravitino job run id - * @return SELECT statement + * @return SELECT statement (one job may cover multiple tables) */ public String selectByJobId(@Param("jobId") long jobId) { return selectColumns() + " WHERE job_id = #{jobId}"; } + /** + * @param jobId Gravitino job run id + * @param tableId table entity id + * @return SELECT statement + */ + public String selectByJobIdAndTableId( + @Param("jobId") long jobId, @Param("tableId") long tableId) { + return selectColumns() + " WHERE job_id = #{jobId} AND table_id = #{tableId}"; + } + + /** + * @param jobId Gravitino job run id + * @param tableId table entity id + * @param policyId policy entity id + * @return SELECT statement + */ + public String selectByJobIdAndTableIdAndPolicyId( + @Param("jobId") long jobId, + @Param("tableId") long tableId, + @Param("policyId") long policyId) { + return selectColumns() + + " WHERE job_id = #{jobId} AND table_id = #{tableId} AND policy_id = #{policyId}"; + } + /** * Sets {@code before_snapshot_id} when {@code S0} is sampled ({@code IS NULL} first write). * Callers treat an already-matching value as an idempotent no-op in Java. * * @param jobId Gravitino job run id + * @param tableId table entity id + * @param policyId policy entity id * @param beforeSnapshotId sampled snapshot id * @return UPDATE statement */ public String updateBeforeSnapshot( - @Param("jobId") long jobId, @Param("beforeSnapshotId") long beforeSnapshotId) { + @Param("jobId") long jobId, + @Param("tableId") long tableId, + @Param("policyId") long policyId, + @Param("beforeSnapshotId") long beforeSnapshotId) { return "UPDATE " + TABLE_NAME + " SET before_snapshot_id = #{beforeSnapshotId}" - + " WHERE job_id = #{jobId} AND before_snapshot_id IS NULL AND finished_at = 0"; + + " WHERE job_id = #{jobId} AND table_id = #{tableId} AND policy_id = #{policyId}" + + " AND before_snapshot_id IS NULL AND finished_at = 0"; } /** - * Sets terminal snapshot pointers when the Spark Job finishes and releases the soft slot. + * Sets terminal snapshot pointers for one bound table/policy row and releases its soft slot. * * <p>Only one in-flight row exists per {@code (table_id, policy_id)} ({@code finished_at = 0}), * so finish timestamps do not contend across concurrent submitters. * * @param jobId Gravitino job run id + * @param tableId table entity id + * @param policyId policy entity id * @param afterSnapshotId terminal snapshot id * @param finishedAt finish time millis; must be {@code > 0} * @return UPDATE statement */ public String markFinished( @Param("jobId") long jobId, + @Param("tableId") long tableId, + @Param("policyId") long policyId, @Param("afterSnapshotId") Long afterSnapshotId, @Param("finishedAt") long finishedAt) { return "UPDATE " + TABLE_NAME + " SET after_snapshot_id = #{afterSnapshotId}, finished_at = #{finishedAt}" - + " WHERE job_id = #{jobId} AND finished_at = 0"; + + " WHERE job_id = #{jobId} AND table_id = #{tableId} AND policy_id = #{policyId}" + + " AND finished_at = 0"; } /** @@ -187,15 +233,21 @@ public class TableMaintenanceActionBaseSQLProvider { * still null (finish listener raced after reclaim). * * @param jobId Gravitino job run id + * @param tableId table entity id + * @param policyId policy entity id * @param afterSnapshotId terminal snapshot id; must be non-null * @return UPDATE statement */ public String backfillAfterSnapshot( - @Param("jobId") long jobId, @Param("afterSnapshotId") long afterSnapshotId) { + @Param("jobId") long jobId, + @Param("tableId") long tableId, + @Param("policyId") long policyId, + @Param("afterSnapshotId") long afterSnapshotId) { return "UPDATE " + TABLE_NAME + " SET after_snapshot_id = #{afterSnapshotId}" - + " WHERE job_id = #{jobId} AND finished_at > 0 AND after_snapshot_id IS NULL"; + + " WHERE job_id = #{jobId} AND table_id = #{tableId} AND policy_id = #{policyId}" + + " AND finished_at > 0 AND after_snapshot_id IS NULL"; } /** diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableMaintenanceActionMetaService.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableMaintenanceActionMetaService.java index 28d47794ec..454243bb10 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableMaintenanceActionMetaService.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableMaintenanceActionMetaService.java @@ -25,6 +25,7 @@ import java.sql.SQLIntegrityConstraintViolationException; import java.util.Collections; import java.util.List; import java.util.Locale; +import java.util.Objects; import java.util.Optional; import javax.annotation.Nullable; import org.apache.gravitino.maintenance.TableMaintenanceConfig; @@ -157,6 +158,10 @@ public class TableMaintenanceActionMetaService { * created_at} from the row: after another node reclaims, that would observe the new fence and * defeat reclaim fencing. The CAS matches zero rows for a stale claim. * + * <p>Idempotent when the row is already bound to the same {@code jobId} and still in-flight + * (Spark may lazy-bind before submit-side bind). Still fails when the claim was reclaimed or + * bound to a different job. + * * <p>This CAS does <em>not</em> prevent duplicate Spark submissions if {@code runJob} already * ran; see {@link #tryAcquireInFlight} submit contract. * @@ -172,10 +177,54 @@ public class TableMaintenanceActionMetaService { SessionUtils.doWithCommitAndFetchResult( TableMaintenanceActionMapper.class, mapper -> mapper.bindJobId(claim.id(), jobId, claim.createdAt())); - Preconditions.checkState( - updated == 1, - "expected to bind job_id on exactly one in-flight action, updated=%s", - updated); + if (updated == 1) { + return; + } + TableMaintenanceActionPO current = getById(claim.id()); + if (current != null + && Objects.equals(current.jobId(), jobId) + && current.finishedAt() != null + && current.finishedAt() == 0L) { + return; + } + throw new IllegalStateException( + String.format( + "expected to bind job_id on exactly one in-flight action, updated=%s", updated)); + } + + /** + * Lazy-binds {@code jobId} onto the unbound in-flight claim for {@code (tableId, policyId)}. + * + * <p>Used when a Spark metrics/finish callback arrives before submit-side {@link #bindJobId}. + * Does nothing (returns {@code false}) when there is no unbound row for that pair. + * + * @param tableId table entity id from the callback + * @param policyId policy entity id from the callback + * @param jobId Gravitino job run id from the callback path + * @return {@code true} when a row is now bound to {@code jobId} + */ + @Monitored( + metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME, + baseMetricName = "bindTableMaintenanceActionJobIdIfUnbound") + public boolean bindJobIdIfUnbound(long tableId, long policyId, long jobId) { + TableMaintenanceActionPO already = getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyId); + if (already != null && already.finishedAt() != null && already.finishedAt() == 0L) { + return true; + } + TableMaintenanceActionPO unbound = + SessionUtils.getWithoutCommit( + TableMaintenanceActionMapper.class, + mapper -> mapper.selectUnboundByTableAndPolicy(tableId, policyId)); + if (unbound == null || unbound.id() == null || unbound.createdAt() == null) { + return false; + } + try { + bindJobId(new OccupancyClaim(unbound.id(), unbound.createdAt()), jobId); + return true; + } catch (IllegalStateException e) { + TableMaintenanceActionPO bound = getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyId); + return bound != null && bound.finishedAt() != null && bound.finishedAt() == 0L; + } } private static OccupancyClaim readUnboundClaim(long tableId, long policyId) { @@ -203,43 +252,108 @@ public class TableMaintenanceActionMetaService { } /** + * Lists all action rows bound to a Gravitino job run. One Spark job may cover multiple tables, so + * the same {@code job_id} can appear on more than one row. + * * @param jobId Gravitino job run id - * @return action row or null + * @return bound action rows (empty when none) + */ + @Monitored( + metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME, + baseMetricName = "listTableMaintenanceActionsByJobId") + public List<TableMaintenanceActionPO> listByJobId(long jobId) { + List<TableMaintenanceActionPO> rows = + SessionUtils.getWithoutCommit( + TableMaintenanceActionMapper.class, mapper -> mapper.selectByJobId(jobId)); + return rows == null ? Collections.emptyList() : rows; + } + + /** + * @param jobId Gravitino job run id + * @param tableId table entity id + * @return action row for that job/table or null + */ + @Nullable + @Monitored( + metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME, + baseMetricName = "getTableMaintenanceActionByJobIdAndTableId") + public TableMaintenanceActionPO getByJobIdAndTableId(long jobId, long tableId) { + return SessionUtils.getWithoutCommit( + TableMaintenanceActionMapper.class, + mapper -> mapper.selectByJobIdAndTableId(jobId, tableId)); + } + + /** + * @param jobId Gravitino job run id + * @param tableId table entity id + * @param policyId policy entity id + * @return action row for that job/table/policy or null + */ + @Nullable + @Monitored( + metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME, + baseMetricName = "getTableMaintenanceActionByJobIdAndTableIdAndPolicyId") + public TableMaintenanceActionPO getByJobIdAndTableIdAndPolicyId( + long jobId, long tableId, long policyId) { + return SessionUtils.getWithoutCommit( + TableMaintenanceActionMapper.class, + mapper -> mapper.selectByJobIdAndTableIdAndPolicyId(jobId, tableId, policyId)); + } + + /** + * Convenience for single-table jobs. Prefer {@link #getByJobIdAndTableIdAndPolicyId}, {@link + * #getByJobIdAndTableId}, or {@link #listByJobId} when a job may cover multiple tables or + * policies. + * + * @param jobId Gravitino job run id + * @return the sole action row, or null when none + * @throws IllegalStateException when more than one table is bound to {@code jobId} */ @Nullable @Monitored( metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME, baseMetricName = "getTableMaintenanceActionByJobId") public TableMaintenanceActionPO getByJobId(long jobId) { - return SessionUtils.getWithoutCommit( - TableMaintenanceActionMapper.class, mapper -> mapper.selectByJobId(jobId)); + List<TableMaintenanceActionPO> rows = listByJobId(jobId); + if (rows.isEmpty()) { + return null; + } + Preconditions.checkState( + rows.size() == 1, + "multiple maintenance actions for job_id=%s; use getByJobIdAndTableIdAndPolicyId or listByJobId", + jobId); + return rows.get(0); } /** - * Records the sampled Iceberg snapshot id ({@code S0}) for a running action. + * Records the sampled Iceberg snapshot id ({@code S0}) for a running action on one table/policy. * * <p>Safe to retry with the same {@code beforeSnapshotId} (cron / in-job re-entry after a partial * sample). A different snapshot id on an already-set row fails. * * @param jobId Gravitino job run id + * @param tableId table entity id + * @param policyId policy entity id * @param beforeSnapshotId sampled snapshot id */ @Monitored( metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME, baseMetricName = "recordTableMaintenanceActionBeforeSnapshot") - public void recordBeforeSnapshot(long jobId, long beforeSnapshotId) { + public void recordBeforeSnapshot(long jobId, long tableId, long policyId, long beforeSnapshotId) { int updated = SessionUtils.doWithCommitAndFetchResult( TableMaintenanceActionMapper.class, - mapper -> mapper.updateBeforeSnapshot(jobId, beforeSnapshotId)); + mapper -> mapper.updateBeforeSnapshot(jobId, tableId, policyId, beforeSnapshotId)); if (updated == 1) { return; } - TableMaintenanceActionPO current = getByJobId(jobId); + TableMaintenanceActionPO current = getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyId); Preconditions.checkState( current != null && current.finishedAt() != null && current.finishedAt() == 0L, - "expected an unfinished action for job_id=%s", - jobId); + "expected an unfinished action for job_id=%s table_id=%s policy_id=%s", + jobId, + tableId, + policyId); Preconditions.checkState( current.beforeSnapshotId() != null && current.beforeSnapshotId() == beforeSnapshotId, "before_snapshot_id already set to %s, refusing %s", @@ -248,29 +362,32 @@ public class TableMaintenanceActionMetaService { } /** - * Marks the action finished with the terminal snapshot id. + * Marks one in-flight action finished for {@code (jobId, tableId, policyId)}. * * <p>{@code UNIQUE(table_id, policy_id, finished_at)} exists so {@code finished_at = 0} is a * single soft slot per pair (like {@code deleted_at}). Only one in-flight row can finish that * pair at a time, so two completed rows cannot race to the same millisecond {@code finished_at}. * - * <p>If failsafe reclaim already closed the row ({@code finished_at > 0}) with a null after - * snapshot, a later finish callback may still backfill a real S1 when {@code afterSnapshotId} is - * non-null. + * <p>If the row was already closed with a null after snapshot, a later call may still backfill a + * real S1 when {@code afterSnapshotId} is non-null. Safe to retry when the row is already + * finished with a null after and {@code afterSnapshotId} is also null. * * @param jobId Gravitino job run id - * @param afterSnapshotId terminal snapshot id + * @param tableId table entity id + * @param policyId policy entity id + * @param afterSnapshotId terminal snapshot id, or null when releasing without an S1 * @param finishedAt finish timestamp millis */ @Monitored( metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME, baseMetricName = "markTableMaintenanceActionFinished") - public void markFinished(long jobId, Long afterSnapshotId, long finishedAt) { + public void markFinished( + long jobId, long tableId, long policyId, Long afterSnapshotId, long finishedAt) { Preconditions.checkArgument(finishedAt > 0, "finishedAt must be positive"); int updated = SessionUtils.doWithCommitAndFetchResult( TableMaintenanceActionMapper.class, - mapper -> mapper.markFinished(jobId, afterSnapshotId, finishedAt)); + mapper -> mapper.markFinished(jobId, tableId, policyId, afterSnapshotId, finishedAt)); if (updated == 1) { return; } @@ -278,19 +395,24 @@ public class TableMaintenanceActionMetaService { int backfilled = SessionUtils.doWithCommitAndFetchResult( TableMaintenanceActionMapper.class, - mapper -> mapper.backfillAfterSnapshot(jobId, afterSnapshotId)); + mapper -> mapper.backfillAfterSnapshot(jobId, tableId, policyId, afterSnapshotId)); Preconditions.checkState( backfilled == 1, - "expected to finish or backfill exactly one action for job_id=%s, finished=%s, backfilled=%s", + "expected to finish or backfill action for job_id=%s table_id=%s policy_id=%s, finished=%s, backfilled=%s", jobId, + tableId, + policyId, updated, backfilled); return; } + TableMaintenanceActionPO current = getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyId); Preconditions.checkState( - false, - "expected to finish exactly one in-flight action for job_id=%s, updated=%s", + current != null && current.finishedAt() != null && current.finishedAt() > 0L, + "expected finished action for job_id=%s table_id=%s policy_id=%s, updated=%s", jobId, + tableId, + policyId, updated); } diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/TableMaintenanceRESTUtils.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/TableMaintenanceRESTUtils.java new file mode 100644 index 0000000000..669912deaa --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/TableMaintenanceRESTUtils.java @@ -0,0 +1,184 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.maintenance.web; + +import java.security.PrivilegedExceptionAction; +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.core.MediaType; +import javax.ws.rs.core.Response; +import org.apache.gravitino.UserPrincipal; +import org.apache.gravitino.auth.AuthConstants; +import org.apache.gravitino.dto.responses.ErrorResponse; +import org.apache.gravitino.exceptions.AlreadyExistsException; +import org.apache.gravitino.exceptions.ForbiddenException; +import org.apache.gravitino.exceptions.NonEmptyEntityException; +import org.apache.gravitino.exceptions.NotFoundException; +import org.apache.gravitino.server.web.ServerHealth; +import org.apache.gravitino.utils.PrincipalUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** REST helpers for table-maintenance resources in the TMS plugin. */ +public final class TableMaintenanceRESTUtils { + + private static final Logger LOG = LoggerFactory.getLogger(TableMaintenanceRESTUtils.class); + + private static final MediaType JSON = MediaType.APPLICATION_JSON_TYPE; + + private TableMaintenanceRESTUtils() {} + + /** + * Returns HTTP 200 with a JSON entity. + * + * @param entity response body + * @param <T> entity type + * @return OK response + */ + public static <T> Response ok(T entity) { + return json(Response.Status.OK, entity); + } + + /** + * Runs {@code action} as the authenticated principal. + * + * @param httpRequest current request + * @param action work to run + * @return the action result + * @throws Exception when the action fails + */ + public static Response doAs( + HttpServletRequest httpRequest, PrivilegedExceptionAction<Response> action) throws Exception { + UserPrincipal principal = + (UserPrincipal) + httpRequest.getAttribute(AuthConstants.AUTHENTICATED_PRINCIPAL_ATTRIBUTE_NAME); + return PrincipalUtils.doAs( + principal == null ? new UserPrincipal(AuthConstants.ANONYMOUS_USER) : principal, action); + } + + /** + * Runs {@code action} and maps failures to the same status codes as the built-in IdP REST utils. + * + * @param httpRequest current request + * @param action work to run + * @param operation short label for logs (e.g. {@code report-before-metrics}) + * @param jobId Gravitino job run id + * @param metalake metalake name + * @return success or error response + */ + public static Response doAs( + HttpServletRequest httpRequest, + PrivilegedExceptionAction<Response> action, + String operation, + long jobId, + String metalake) { + try { + return doAs(httpRequest, action); + } catch (Exception e) { + return handleException(operation, jobId, metalake, e); + } + } + + /** + * Maps a maintenance-action failure to an error response. + * + * @param operation short label for logs + * @param jobId Gravitino job run id + * @param metalake metalake name + * @param e failure + * @return error response + */ + public static Response handleException( + String operation, long jobId, String metalake, Exception e) { + ServerHealth.getInstance().recordFailure(e); + String errorMsg = + String.format( + "Failed to operate maintenance action [%s] operation [%s] under metalake [%s], " + + "reason [%s]", + jobId, operation, metalake, e.getMessage()); + LOG.warn(errorMsg, e); + return toErrorResponse(errorMsg, e); + } + + private static Response toErrorResponse(String errorMsg, Exception e) { + if (e instanceof IllegalArgumentException) { + return illegalArguments(errorMsg, e); + } + if (e instanceof ForbiddenException) { + return forbidden(errorMsg, e); + } + if (e instanceof NotFoundException) { + return notFound(errorMsg, e); + } + if (e instanceof AlreadyExistsException) { + return alreadyExists(errorMsg, e); + } + if (e instanceof NonEmptyEntityException) { + return json( + Response.Status.CONFLICT, + ErrorResponse.nonEmpty(e.getClass().getSimpleName(), errorMsg, e)); + } + if (e instanceof UnsupportedOperationException) { + return unsupportedOperation(errorMsg, e); + } + return internalError(errorMsg, e); + } + + private static Response illegalArguments(String message, Throwable throwable) { + return json( + Response.Status.BAD_REQUEST, + ErrorResponse.illegalArguments( + exceptionType(throwable, "IllegalArgumentException"), message, throwable)); + } + + private static Response notFound(String message, Throwable throwable) { + return json( + Response.Status.NOT_FOUND, + ErrorResponse.notFound(exceptionType(throwable, "NotFoundException"), message, throwable)); + } + + private static Response alreadyExists(String message, Throwable throwable) { + return json( + Response.Status.CONFLICT, + ErrorResponse.alreadyExists( + exceptionType(throwable, "AlreadyExistsException"), message, throwable)); + } + + private static Response forbidden(String message, Throwable throwable) { + return json(Response.Status.FORBIDDEN, ErrorResponse.forbidden(message, throwable)); + } + + private static Response unsupportedOperation(String message, Throwable throwable) { + return json( + Response.Status.NOT_IMPLEMENTED, ErrorResponse.unsupportedOperation(message, throwable)); + } + + private static Response internalError(String message, Throwable throwable) { + ServerHealth.getInstance().recordFailure(throwable); + return json( + Response.Status.INTERNAL_SERVER_ERROR, ErrorResponse.internalError(message, throwable)); + } + + private static Response json(Response.Status status, Object entity) { + return Response.status(status).entity(entity).type(JSON).build(); + } + + private static String exceptionType(Throwable throwable, String defaultType) { + return throwable == null ? defaultType : throwable.getClass().getSimpleName(); + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/TableMaintenanceActionOperations.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/TableMaintenanceActionOperations.java new file mode 100644 index 0000000000..a2130e0267 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/TableMaintenanceActionOperations.java @@ -0,0 +1,188 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.maintenance.web.rest; + +import com.codahale.metrics.annotation.ResponseMetered; +import com.codahale.metrics.annotation.Timed; +import com.google.common.annotations.VisibleForTesting; +import javax.inject.Inject; +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.Consumes; +import javax.ws.rs.DefaultValue; +import javax.ws.rs.POST; +import javax.ws.rs.Path; +import javax.ws.rs.PathParam; +import javax.ws.rs.Produces; +import javax.ws.rs.QueryParam; +import javax.ws.rs.core.Context; +import javax.ws.rs.core.MediaType; +import javax.ws.rs.core.Response; +import org.apache.gravitino.Entity; +import org.apache.gravitino.dto.responses.BaseResponse; +import org.apache.gravitino.maintenance.MaintenanceMetricsService; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionFinishRequest; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionMetricsRequest; +import org.apache.gravitino.maintenance.web.TableMaintenanceRESTUtils; +import org.apache.gravitino.metrics.MetricNames; +import org.apache.gravitino.rest.RESTRequest; +import org.apache.gravitino.server.authorization.annotations.AuthorizationExpression; +import org.apache.gravitino.server.authorization.annotations.AuthorizationMetadata; + +/** + * Spark callbacks for before/after snapshot metrics and soft-slot finish on an in-flight + * maintenance action. + * + * <p>Continue / stop stays in the job. Soft-slot release is via {@code finish=true} query param on + * metrics, or {@code POST .../finish} when no metrics are written. + */ +@Path("metalakes/{metalake}/maintenance/actions") +@Consumes({ + MediaType.APPLICATION_JSON, + MediaType.APPLICATION_JSON + "; charset=UTF-8", + "application/vnd.gravitino.v1+json", + "application/vnd.gravitino.v1+json; charset=UTF-8" +}) +@Produces("application/vnd.gravitino.v1+json") +public class TableMaintenanceActionOperations { + + private final MaintenanceMetricsService metricsService; + + @Context private HttpServletRequest httpRequest; + + /** + * Creates the resource. + * + * @param metricsService metrics report domain service + */ + @Inject + public TableMaintenanceActionOperations(MaintenanceMetricsService metricsService) { + this.metricsService = metricsService; + } + + /** + * Persists S0 metrics and {@code before_snapshot_id} for the action bound to {@code jobId}. + * + * @param metalake metalake name from the request path + * @param jobId Gravitino job run id bound to the action + * @param finish when {@code true}, close the soft slot after storing metrics + * @param request metrics payload + * @return empty success response + */ + @POST + @Path("{jobId}/metrics/before") + @Timed(name = "tms-action-metrics-before." + MetricNames.HTTP_PROCESS_DURATION, absolute = true) + @ResponseMetered(name = "tms-action-metrics-before", absolute = true) + @AuthorizationExpression(expression = "METALAKE::OWNER || METALAKE::RUN_JOB") + public Response reportBeforeMetrics( + @PathParam("metalake") @AuthorizationMetadata(type = Entity.EntityType.METALAKE) + String metalake, + @PathParam("jobId") long jobId, + @QueryParam("finish") @DefaultValue("false") boolean finish, + MaintenanceActionMetricsRequest request) { + return TableMaintenanceRESTUtils.doAs( + httpRequest, + () -> { + requireBody(request); + metricsService.reportBeforeMetrics(metalake, jobId, finish, request); + return TableMaintenanceRESTUtils.ok(new BaseResponse()); + }, + "report-before-metrics", + jobId, + metalake); + } + + /** + * Persists S1 metrics for the action bound to {@code jobId}. + * + * @param metalake metalake name from the request path + * @param jobId Gravitino job run id bound to the action + * @param finish when {@code true}, close the soft slot after storing metrics + * @param request metrics payload + * @return empty success response + */ + @POST + @Path("{jobId}/metrics/after") + @Timed(name = "tms-action-metrics-after." + MetricNames.HTTP_PROCESS_DURATION, absolute = true) + @ResponseMetered(name = "tms-action-metrics-after", absolute = true) + @AuthorizationExpression(expression = "METALAKE::OWNER || METALAKE::RUN_JOB") + public Response reportAfterMetrics( + @PathParam("metalake") @AuthorizationMetadata(type = Entity.EntityType.METALAKE) + String metalake, + @PathParam("jobId") long jobId, + @QueryParam("finish") @DefaultValue("false") boolean finish, + MaintenanceActionMetricsRequest request) { + return TableMaintenanceRESTUtils.doAs( + httpRequest, + () -> { + requireBody(request); + metricsService.reportAfterMetrics(metalake, jobId, finish, request); + return TableMaintenanceRESTUtils.ok(new BaseResponse()); + }, + "report-after-metrics", + jobId, + metalake); + } + + /** + * Closes the soft slot for one bound table without writing metrics. + * + * @param metalake metalake name from the request path + * @param jobId Gravitino job run id bound to the action + * @param request finish payload + * @return empty success response + */ + @POST + @Path("{jobId}/finish") + @Timed(name = "tms-action-finish." + MetricNames.HTTP_PROCESS_DURATION, absolute = true) + @ResponseMetered(name = "tms-action-finish", absolute = true) + @AuthorizationExpression(expression = "METALAKE::OWNER || METALAKE::RUN_JOB") + public Response finishAction( + @PathParam("metalake") @AuthorizationMetadata(type = Entity.EntityType.METALAKE) + String metalake, + @PathParam("jobId") long jobId, + MaintenanceActionFinishRequest request) { + return TableMaintenanceRESTUtils.doAs( + httpRequest, + () -> { + requireBody(request); + metricsService.finishAction(metalake, jobId, request); + return TableMaintenanceRESTUtils.ok(new BaseResponse()); + }, + "finish", + jobId, + metalake); + } + + private static void requireBody(RESTRequest request) { + if (request == null) { + throw new IllegalArgumentException("Request body cannot be null"); + } + request.validate(); + } + + /** + * Installs the request used by {@link TableMaintenanceRESTUtils#doAs}. + * + * @param httpRequest request whose authenticated principal is adopted + */ + @VisibleForTesting + void setHttpRequest(HttpServletRequest httpRequest) { + this.httpRequest = httpRequest; + } +} diff --git a/docs/build.gradle.kts b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/TableMaintenanceBinder.java similarity index 53% copy from docs/build.gradle.kts copy to plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/TableMaintenanceBinder.java index 39810c3dda..61edd478ff 100644 --- a/docs/build.gradle.kts +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/TableMaintenanceBinder.java @@ -16,35 +16,16 @@ * specific language governing permissions and limitations * under the License. */ +package org.apache.gravitino.maintenance.web.rest; -import com.github.gradle.node.NodeExtension -import com.github.gradle.node.npm.task.NpxTask +import org.apache.gravitino.maintenance.MaintenanceMetricsService; +import org.glassfish.hk2.utilities.binding.AbstractBinder; -configure<NodeExtension> { - version = "21.6.1" - npmVersion = "10.2.4" - download = true -} - -tasks { - val lintOpenAPI by registering(NpxTask::class) { - command.set("@redocly/[email protected]") - args.set(listOf("lint", "--extends=recommended-strict", "${project.projectDir}/open-api/openapi.yaml")) - } - - val lintIdpOpenAPI by registering(NpxTask::class) { - command.set("@redocly/[email protected]") - args.set( - listOf( - "lint", - "--extends=recommended-strict", - "${project.projectDir}/open-api/idp/openapi.yaml" - ) - ) - } +/** Registers TMS maintenance metrics REST dependencies for Jersey/HK2 injection. */ +public class TableMaintenanceBinder extends AbstractBinder { - build { - dependsOn(lintOpenAPI) - dependsOn(lintIdpOpenAPI) + @Override + protected void configure() { + bind(MaintenanceMetricsService.create()).to(MaintenanceMetricsService.class); } } diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/feature/TableMaintenanceRESTFeature.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/feature/TableMaintenanceRESTFeature.java index 3b3f7320ee..1b9473553b 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/feature/TableMaintenanceRESTFeature.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/feature/TableMaintenanceRESTFeature.java @@ -31,6 +31,8 @@ import org.apache.gravitino.maintenance.policy.PolicyExpandLifecycleHook; import org.apache.gravitino.maintenance.policy.PolicyExpandReconciler; import org.apache.gravitino.maintenance.policy.PolicyExpandRegistrar; import org.apache.gravitino.maintenance.scheduler.TableMaintenanceSchedulerPools; +import org.apache.gravitino.maintenance.web.rest.TableMaintenanceActionOperations; +import org.apache.gravitino.maintenance.web.rest.TableMaintenanceBinder; import org.apache.gravitino.policy.hook.PolicyLifecycleHooks; import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper; import org.glassfish.jersey.server.monitoring.ApplicationEvent; @@ -41,16 +43,16 @@ import org.slf4j.Logger; import org.slf4j.LoggerFactory; /** - * Jersey feature that starts table maintenance db-scheduler pools when the extension package is - * enabled (§7.1). + * Jersey feature that starts table maintenance db-scheduler pools and registers Spark action + * callback REST when the extension package is enabled (§7.1 / §5.2.2). * - * <p>TMS does not register public REST resources in this change; the feature only owns server - * lifecycle for in-process scheduling. Cron and commit pools poll. The expand pool stays stopped, - * and a policy lifecycle hook writes the matching {@code policy-expand} row on the same transaction - * as create, alter, enable, disable, and delete. On start, existing maintenance policies are - * reconciled into expand rows so upgrade and first-enable paths do not wait for an edit. Pools stop - * on Jersey {@code DESTROY_FINISHED}, which runs during Jetty shutdown before {@code GravitinoEnv} - * closes the entity-store {@code DataSource}. + * <p>Registers {@link TableMaintenanceActionOperations} for sample/finish callbacks. Cron and + * commit pools poll. The expand pool stays stopped, and a policy lifecycle hook writes the matching + * {@code policy-expand} row on the same transaction as create, alter, enable, disable, and delete. + * On start, existing maintenance policies are reconciled into expand rows so upgrade and + * first-enable paths do not wait for an edit. Pools stop on Jersey {@code DESTROY_FINISHED}, which + * runs during Jetty shutdown before {@code GravitinoEnv} closes the entity-store {@code + * DataSource}. * * <p>After {@link #stopSchedulers()}, a later {@link #configure} may start pools again, re-register * the destroy listener, and reconcile again (needed for tests / embedded restarts in the same JVM). @@ -104,6 +106,8 @@ public class TableMaintenanceRESTFeature implements Feature { // Register only after a successful CAS so a later configure-after-stop can attach a listener // again on the new FeatureContext (listener is scoped to the Jersey Application). context.register(new SchedulerShutdownListener()); + context.register(TableMaintenanceBinder.class); + context.register(TableMaintenanceActionOperations.class); LOG.info("Table maintenance scheduler pools are running"); return true; } diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/occupancy/TestOccupancyGate.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/occupancy/TestOccupancyGate.java index 1479343efd..d522652fd6 100644 --- a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/occupancy/TestOccupancyGate.java +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/occupancy/TestOccupancyGate.java @@ -18,41 +18,30 @@ */ package org.apache.gravitino.maintenance.occupancy; -import static org.apache.gravitino.maintenance.storage.TableMaintenanceActionTestUtils.expireClaim; - import com.google.common.collect.ImmutableMap; import com.google.common.collect.ImmutableSet; import java.io.IOException; -import java.util.function.Function; import org.apache.gravitino.MetadataObject; import org.apache.gravitino.Namespace; -import org.apache.gravitino.job.JobHandle; import org.apache.gravitino.maintenance.TableMaintenanceConfig; import org.apache.gravitino.maintenance.storage.service.OccupancyClaim; import org.apache.gravitino.maintenance.storage.service.TableMaintenanceActionMetaService; -import org.apache.gravitino.meta.JobEntity; -import org.apache.gravitino.meta.JobTemplateEntity; import org.apache.gravitino.meta.PolicyEntity; import org.apache.gravitino.meta.TableEntity; import org.apache.gravitino.policy.PolicyContents; import org.apache.gravitino.storage.relational.TestJDBCBackend; -import org.apache.gravitino.storage.relational.service.JobMetaService; -import org.apache.gravitino.utils.NameIdentifierUtil; -import org.apache.gravitino.utils.NamespaceUtil; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.TestTemplate; class TestOccupancyGate extends TestJDBCBackend { @TestTemplate - void testInFlightReclaimsTerminalBoundJobFinishedAt() throws IOException { + void testUnboundClaimAndBoundSlotOccupyUntilMarkFinished() throws IOException { String metalakeName = "tms_occ_ml"; createAndInsertMakeLake(metalakeName); createAndInsertCatalog(metalakeName, "cat"); createAndInsertSchema(metalakeName, "cat", "sch"); TableEntity table = createAndInsertTableEntity(Namespace.of(metalakeName, "cat", "sch"), "tbl"); - JobTemplateEntity template = - createAndInsertShellJobTemplateEntity("occ-template", "comment", metalakeName); TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); OccupancyGate gate = new OccupancyGate(actions); @@ -65,52 +54,20 @@ class TestOccupancyGate extends TestJDBCBackend { table.id(), policyId, TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS); Assertions.assertTrue(gate.isInFlight(table.id(), policyId)); - JobEntity started = - JobEntity.builder() - .withId(7001L) - .withJobExecutionId("exec-7001") - .withNamespace(NamespaceUtil.ofJob(metalakeName)) - .withJobTemplateName(template.name()) - .withStatus(JobHandle.Status.STARTED) - .withAuditInfo(AUDIT_INFO) - .withStartedAt(System.currentTimeMillis()) - .withFinishedAt(0L) - .build(); - backend.insert(started, false); - actions.bindJobId(claim, started.id()); - + actions.bindJobId(claim, 7001L); Assertions.assertTrue(gate.isInFlight(table.id(), policyId)); Assertions.assertTrue( gate.shouldSkip(table.id(), policyId, 3_600_000L, System.currentTimeMillis())); - Assertions.assertEquals(0L, actions.getByJobId(started.id()).finishedAt()); - - long finishedAt = 1_700_000_000_000L; - JobMetaService.getInstance() - .updateJob( - NameIdentifierUtil.ofJob(metalakeName, started.name()), - (Function<JobEntity, JobEntity>) - current -> - JobEntity.builder() - .withId(current.id()) - .withJobExecutionId(current.jobExecutionId()) - .withNamespace(current.namespace()) - .withJobTemplateName(current.jobTemplateName()) - .withStatus(JobHandle.Status.SUCCEEDED) - .withAuditInfo(current.auditInfo()) - .withStartedAt(current.startedAt()) - .withFinishedAt(finishedAt) - .build()); + long finishedAt = System.currentTimeMillis(); + actions.markFinished(7001L, table.id(), policyId, 200L, finishedAt); Assertions.assertFalse(gate.isInFlight(table.id(), policyId)); - Assertions.assertEquals(finishedAt, actions.getByJobId(started.id()).finishedAt()); - Assertions.assertTrue( - gate.withinMinInterval(table.id(), policyId, 3_600_000L, finishedAt + 1_000L)); - Assertions.assertFalse( - gate.withinMinInterval(table.id(), policyId, 3_600_000L, finishedAt + 3_600_001L)); + Assertions.assertEquals( + finishedAt, actions.getByJobIdAndTableId(7001L, table.id()).finishedAt()); } @TestTemplate - void testMinIntervalUsesMarkedFinishedAt() throws IOException { + void testMinIntervalUsesActionFinishedAtOnly() throws IOException { String metalakeName = "tms_occ_ml2"; createAndInsertMakeLake(metalakeName); createAndInsertCatalog(metalakeName, "cat"); @@ -119,242 +76,39 @@ class TestOccupancyGate extends TestJDBCBackend { TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); OccupancyGate gate = new OccupancyGate(actions); - long policyId = 44L; - long jobId = 7002L; - long finishedAt = 1_700_000_100_000L; + long policyId = insertPolicy(metalakeName, "occ-policy-2").id(); OccupancyClaim claim = actions.insertInFlight( table.id(), policyId, TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS); + long jobId = 8001L; actions.bindJobId(claim, jobId); - actions.markFinished(jobId, 200L, finishedAt); - Assertions.assertFalse(gate.isInFlight(table.id(), policyId)); + long finishedAt = 1_700_000_000_000L; + actions.markFinished(jobId, table.id(), policyId, null, finishedAt); + Assertions.assertTrue( - gate.withinMinInterval(table.id(), policyId, 3_600_000L, finishedAt + 1_000L)); + gate.withinMinInterval(table.id(), policyId, 60_000L, finishedAt + 1_000L)); + Assertions.assertFalse( + gate.withinMinInterval(table.id(), policyId, 60_000L, finishedAt + 120_000L)); } @TestTemplate - void testTerminalBoundJobIsReclaimedSoSlotCanBeAcquiredAgain() throws IOException { + void testExpiredUnboundClaimDoesNotBlock() throws Exception { String metalakeName = "tms_occ_ml3"; createAndInsertMakeLake(metalakeName); createAndInsertCatalog(metalakeName, "cat"); createAndInsertSchema(metalakeName, "cat", "sch"); TableEntity table = createAndInsertTableEntity(Namespace.of(metalakeName, "cat", "sch"), "tbl"); - JobTemplateEntity template = - createAndInsertShellJobTemplateEntity("occ-template-3", "comment", metalakeName); TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); - OccupancyGate gate = new OccupancyGate(actions); + long timeoutMs = 200L; + OccupancyGate gate = new OccupancyGate(actions, timeoutMs); long policyId = insertPolicy(metalakeName, "occ-policy-3").id(); - long finishedAt = 1_700_000_200_000L; - JobEntity succeeded = - JobEntity.builder() - .withId(7003L) - .withJobExecutionId("exec-7003") - .withNamespace(NamespaceUtil.ofJob(metalakeName)) - .withJobTemplateName(template.name()) - .withStatus(JobHandle.Status.SUCCEEDED) - .withAuditInfo(AUDIT_INFO) - .withStartedAt(finishedAt - 1_000L) - .withFinishedAt(finishedAt) - .build(); - backend.insert(succeeded, false); - OccupancyClaim stale = - actions.insertInFlight( - table.id(), policyId, TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS); - actions.bindJobId(stale, succeeded.id()); - Assertions.assertEquals(0L, actions.getByJobId(succeeded.id()).finishedAt()); - Assertions.assertFalse(gate.isInFlight(table.id(), policyId)); - Assertions.assertEquals(finishedAt, actions.getByJobId(succeeded.id()).finishedAt()); - Assertions.assertEquals(0, actions.listInFlight(table.id(), policyId).size()); - - JobEntity started = - JobEntity.builder() - .withId(7004L) - .withJobExecutionId("exec-7004") - .withNamespace(NamespaceUtil.ofJob(metalakeName)) - .withJobTemplateName(template.name()) - .withStatus(JobHandle.Status.STARTED) - .withAuditInfo(AUDIT_INFO) - .withStartedAt(System.currentTimeMillis()) - .withFinishedAt(0L) - .build(); - backend.insert(started, false); - OccupancyClaim active = - actions.insertInFlight( - table.id(), policyId, TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS); - actions.bindJobId(active, started.id()); - - Assertions.assertEquals(1, actions.listInFlight(table.id(), policyId).size()); + actions.insertInFlight(table.id(), policyId, timeoutMs); Assertions.assertTrue(gate.isInFlight(table.id(), policyId)); - Assertions.assertTrue( - gate.shouldSkip(table.id(), policyId, 3_600_000L, finishedAt + 3_600_001L)); - } - - @TestTemplate - void testMissingJobMetaIsReclaimedSoSlotCanBeAcquiredAgain() throws IOException { - String metalakeName = "tms_occ_ml5"; - createAndInsertMakeLake(metalakeName); - createAndInsertCatalog(metalakeName, "cat"); - createAndInsertSchema(metalakeName, "cat", "sch"); - TableEntity table = createAndInsertTableEntity(Namespace.of(metalakeName, "cat", "sch"), "tbl"); - - TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); - OccupancyGate gate = new OccupancyGate(actions); - // Policy id need not exist: job lookup goes through the table's metalake. - long policyId = 77L; - long orphanJobId = 7005L; - - OccupancyClaim stale = - actions.insertInFlight( - table.id(), policyId, TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS); - actions.bindJobId(stale, orphanJobId); - Assertions.assertEquals(0L, actions.getByJobId(orphanJobId).finishedAt()); - // No JobEntity row for orphanJobId — reclaim must still free the soft slot. + Thread.sleep(timeoutMs + 50L); Assertions.assertFalse(gate.isInFlight(table.id(), policyId)); - Assertions.assertTrue(actions.getByJobId(orphanJobId).finishedAt() > 0L); - Assertions.assertEquals(0, actions.listInFlight(table.id(), policyId).size()); - Assertions.assertTrue( - actions - .tryAcquireInFlight( - table.id(), policyId, TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS) - .isPresent()); - } - - @TestTemplate - void testReclaimDoesNotCopyBeforeSnapshotIdIntoAfter() throws IOException { - String metalakeName = "tms_occ_ml6"; - createAndInsertMakeLake(metalakeName); - createAndInsertCatalog(metalakeName, "cat"); - createAndInsertSchema(metalakeName, "cat", "sch"); - TableEntity table = createAndInsertTableEntity(Namespace.of(metalakeName, "cat", "sch"), "tbl"); - JobTemplateEntity template = - createAndInsertShellJobTemplateEntity("occ-template-6", "comment", metalakeName); - - TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); - OccupancyGate gate = new OccupancyGate(actions); - long policyId = insertPolicy(metalakeName, "occ-policy-6").id(); - long finishedAt = 1_700_000_300_000L; - long beforeSnapshotId = 100L; - - JobEntity succeeded = - JobEntity.builder() - .withId(7006L) - .withJobExecutionId("exec-7006") - .withNamespace(NamespaceUtil.ofJob(metalakeName)) - .withJobTemplateName(template.name()) - .withStatus(JobHandle.Status.SUCCEEDED) - .withAuditInfo(AUDIT_INFO) - .withStartedAt(finishedAt - 1_000L) - .withFinishedAt(finishedAt) - .build(); - backend.insert(succeeded, false); - OccupancyClaim claim = - actions.insertInFlight( - table.id(), policyId, TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS); - actions.bindJobId(claim, succeeded.id()); - actions.recordBeforeSnapshot(succeeded.id(), beforeSnapshotId); - Assertions.assertEquals( - beforeSnapshotId, actions.getByJobId(succeeded.id()).beforeSnapshotId()); - Assertions.assertNull(actions.getByJobId(succeeded.id()).afterSnapshotId()); - - Assertions.assertFalse(gate.isInFlight(table.id(), policyId)); - Assertions.assertEquals(finishedAt, actions.getByJobId(succeeded.id()).finishedAt()); - Assertions.assertEquals( - beforeSnapshotId, actions.getByJobId(succeeded.id()).beforeSnapshotId()); - // Failsafe reclaim must not invent S1 from S0 (§5.2.3). - Assertions.assertNull(actions.getByJobId(succeeded.id()).afterSnapshotId()); - - // Late finish listener may still backfill a real S1 after reclaim closed the soft slot. - long afterSnapshotId = 200L; - actions.markFinished(succeeded.id(), afterSnapshotId, System.currentTimeMillis()); - Assertions.assertEquals(afterSnapshotId, actions.getByJobId(succeeded.id()).afterSnapshotId()); - Assertions.assertEquals(finishedAt, actions.getByJobId(succeeded.id()).finishedAt()); - } - - @TestTemplate - void testActiveJobNotReclaimedWhenPolicyIdIsUnknown() throws IOException { - String metalakeName = "tms_occ_ml7"; - createAndInsertMakeLake(metalakeName); - createAndInsertCatalog(metalakeName, "cat"); - createAndInsertSchema(metalakeName, "cat", "sch"); - TableEntity table = createAndInsertTableEntity(Namespace.of(metalakeName, "cat", "sch"), "tbl"); - JobTemplateEntity template = - createAndInsertShellJobTemplateEntity("occ-template-7", "comment", metalakeName); - - TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); - OccupancyGate gate = new OccupancyGate(actions); - // Simulate a deleted / unreadable policy: lookup must still use the table metalake. - long unknownPolicyId = 9_000_007L; - - JobEntity started = - JobEntity.builder() - .withId(7007L) - .withJobExecutionId("exec-7007") - .withNamespace(NamespaceUtil.ofJob(metalakeName)) - .withJobTemplateName(template.name()) - .withStatus(JobHandle.Status.STARTED) - .withAuditInfo(AUDIT_INFO) - .withStartedAt(System.currentTimeMillis()) - .withFinishedAt(0L) - .build(); - backend.insert(started, false); - OccupancyClaim claim = - actions.insertInFlight( - table.id(), unknownPolicyId, TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS); - actions.bindJobId(claim, started.id()); - - Assertions.assertTrue(gate.isInFlight(table.id(), unknownPolicyId)); - Assertions.assertEquals(0L, actions.getByJobId(started.id()).finishedAt()); - Assertions.assertEquals(1, actions.listInFlight(table.id(), unknownPolicyId).size()); - } - - @TestTemplate - void testUnresolvedTableMetalakeDoesNotReclaimBoundRow() throws IOException { - TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); - OccupancyGate gate = new OccupancyGate(actions); - long missingTableId = 9_999_001L; - long policyId = 88L; - long orphanJobId = 7008L; - - OccupancyClaim claim = - actions.insertInFlight( - missingTableId, policyId, TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS); - actions.bindJobId(claim, orphanJobId); - - // Table/metalake cannot be resolved — do not treat the job as missing/terminal. - Assertions.assertTrue(gate.isInFlight(missingTableId, policyId)); - Assertions.assertEquals(0L, actions.getByJobId(orphanJobId).finishedAt()); - Assertions.assertEquals(1, actions.listInFlight(missingTableId, policyId).size()); - } - - @TestTemplate - void testExpiredUnboundClaimIsNotInFlightOnDbClock() throws Exception { - String metalakeName = "tms_occ_ml4"; - createAndInsertMakeLake(metalakeName); - createAndInsertCatalog(metalakeName, "cat"); - createAndInsertSchema(metalakeName, "cat", "sch"); - TableEntity table = createAndInsertTableEntity(Namespace.of(metalakeName, "cat", "sch"), "tbl"); - - TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); - long claimTimeoutMs = TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS; - OccupancyGate gate = new OccupancyGate(actions, claimTimeoutMs); - long policyId = 66L; - - OccupancyClaim claim = actions.insertInFlight(table.id(), policyId, claimTimeoutMs); - Assertions.assertTrue(actions.hasLiveUnboundClaim(table.id(), policyId, claimTimeoutMs)); - Assertions.assertTrue(gate.isInFlight(table.id(), policyId)); - - expireClaim(claim, claimTimeoutMs); - - // Unbound TTL is evaluated in the DAO against the database clock, so OccupancyGate agrees with - // acquire/reclaim even when the JVM clock would disagree (multi-node skew). - Assertions.assertFalse(actions.hasLiveUnboundClaim(table.id(), policyId, claimTimeoutMs)); - Assertions.assertFalse(gate.isInFlight(table.id(), policyId)); - Assertions.assertTrue( - actions.tryAcquireInFlight(table.id(), policyId, claimTimeoutMs).isPresent()); - Assertions.assertEquals(claim.id(), actions.getById(claim.id()).id()); } private PolicyEntity insertPolicy(String metalakeName, String policyName) throws IOException { diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/storage/service/TestTableMaintenanceActionMetaService.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/storage/service/TestTableMaintenanceActionMetaService.java index 7c66eb4ad5..ec54a63786 100644 --- a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/storage/service/TestTableMaintenanceActionMetaService.java +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/storage/service/TestTableMaintenanceActionMetaService.java @@ -42,38 +42,71 @@ class TestTableMaintenanceActionMetaService extends TestJDBCBackend { long after = System.currentTimeMillis(); Assertions.assertNull(actionService.getById(claim.id()).jobId()); Assertions.assertEquals(0L, actionService.getById(claim.id()).finishedAt()); - // DB clock may skew slightly from the JVM; allow a small window. Assertions.assertTrue(claim.createdAt() >= before - 5_000L); Assertions.assertTrue(claim.createdAt() <= after + 5_000L); Assertions.assertEquals(claim.createdAt(), actionService.getById(claim.id()).createdAt()); long jobId = 999L; + long tableId = 10L; + long policyId = 20L; actionService.bindJobId(claim, jobId); - Assertions.assertEquals(jobId, actionService.getByJobId(jobId).jobId()); - Assertions.assertEquals(0L, actionService.getByJobId(jobId).finishedAt()); - - actionService.recordBeforeSnapshot(jobId, 100L); - Assertions.assertEquals(100L, actionService.getByJobId(jobId).beforeSnapshotId()); - // Same S0 is idempotent (in-job / cron retry after a partial sample). - actionService.recordBeforeSnapshot(jobId, 100L); - Assertions.assertEquals(100L, actionService.getByJobId(jobId).beforeSnapshotId()); + Assertions.assertEquals(jobId, actionService.getByJobIdAndTableId(jobId, tableId).jobId()); + Assertions.assertEquals(0L, actionService.getByJobIdAndTableId(jobId, tableId).finishedAt()); + + actionService.recordBeforeSnapshot(jobId, tableId, policyId, 100L); + Assertions.assertEquals( + 100L, actionService.getByJobIdAndTableId(jobId, tableId).beforeSnapshotId()); + actionService.recordBeforeSnapshot(jobId, tableId, policyId, 100L); + Assertions.assertEquals( + 100L, actionService.getByJobIdAndTableId(jobId, tableId).beforeSnapshotId()); Assertions.assertThrows( - IllegalStateException.class, () -> actionService.recordBeforeSnapshot(jobId, 101L)); + IllegalStateException.class, + () -> actionService.recordBeforeSnapshot(jobId, tableId, policyId, 101L)); - metricsService.insertMetrics(10L, 20L, 100L, "{\"mse\":1}"); - actionService.markFinished(jobId, 200L, System.currentTimeMillis()); + metricsService.insertMetrics(tableId, policyId, 100L, "{\"mse\":1}"); + actionService.markFinished(jobId, tableId, policyId, 200L, System.currentTimeMillis()); Assertions.assertThrows( IllegalStateException.class, - () -> actionService.markFinished(jobId, 201L, System.currentTimeMillis())); - metricsService.insertMetrics(10L, 20L, 200L, "{\"mse\":2}"); + () -> + actionService.markFinished(jobId, tableId, policyId, 201L, System.currentTimeMillis())); + metricsService.insertMetrics(tableId, policyId, 200L, "{\"mse\":2}"); actionService.deleteByJobId(jobId); - Assertions.assertNull(actionService.getByJobId(jobId)); + Assertions.assertTrue(actionService.listByJobId(jobId).isEmpty()); Assertions.assertEquals( - "{\"mse\":1}", metricsService.getMetrics(10L, 20L, 100L).metricsValue()); + "{\"mse\":1}", metricsService.getMetrics(tableId, policyId, 100L).metricsValue()); Assertions.assertEquals( - "{\"mse\":2}", metricsService.getMetrics(10L, 20L, 200L).metricsValue()); + "{\"mse\":2}", metricsService.getMetrics(tableId, policyId, 200L).metricsValue()); + } + + @TestTemplate + void testOneJobBindsMultipleTablesAndMarkFinishedClosesPerRow() { + TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); + TableMaintenanceConfig config = TableMaintenanceConfig.forTest(1, 1, 1, 500L, 1_000L, 4, 10); + long jobId = 777L; + OccupancyClaim claimA = actions.insertInFlight(101L, 201L, config); + OccupancyClaim claimB = actions.insertInFlight(102L, 202L, config); + actions.bindJobId(claimA, jobId); + actions.bindJobId(claimB, jobId); + + Assertions.assertEquals(2, actions.listByJobId(jobId).size()); + Assertions.assertThrows(IllegalStateException.class, () -> actions.getByJobId(jobId)); + + actions.recordBeforeSnapshot(jobId, 101L, 201L, 11L); + actions.recordBeforeSnapshot(jobId, 102L, 202L, 22L); + + long finishedAtA = System.currentTimeMillis(); + actions.markFinished(jobId, 101L, 201L, null, finishedAtA); + Assertions.assertEquals(finishedAtA, actions.getByJobIdAndTableId(jobId, 101L).finishedAt()); + Assertions.assertEquals(0L, actions.getByJobIdAndTableId(jobId, 102L).finishedAt()); + + long finishedAtB = finishedAtA + 10; + actions.markFinished(jobId, 102L, 202L, 33L, finishedAtB); + Assertions.assertEquals(finishedAtB, actions.getByJobIdAndTableId(jobId, 102L).finishedAt()); + Assertions.assertEquals(33L, actions.getByJobIdAndTableId(jobId, 102L).afterSnapshotId()); + // Idempotent null-after finish for already closed A. + actions.markFinished(jobId, 101L, 201L, null, finishedAtA + 1); } @TestTemplate @@ -81,25 +114,123 @@ class TestTableMaintenanceActionMetaService extends TestJDBCBackend { TableMaintenanceActionMetaService actionService = TableMaintenanceActionMetaService.getInstance(); long jobId = 1999L; + long tableId = 12L; + long policyId = 22L; OccupancyClaim claim = actionService.insertInFlight( - 12L, 22L, TableMaintenanceConfig.forTest(1, 1, 1, 500L, 1_000L, 4, 10)); + tableId, policyId, TableMaintenanceConfig.forTest(1, 1, 1, 500L, 1_000L, 4, 10)); actionService.bindJobId(claim, jobId); - actionService.recordBeforeSnapshot(jobId, 100L); + actionService.recordBeforeSnapshot(jobId, tableId, policyId, 100L); long reclaimFinishedAt = System.currentTimeMillis(); - // Failsafe reclaim closes with null S1. - actionService.markFinished(jobId, null, reclaimFinishedAt); - Assertions.assertEquals(reclaimFinishedAt, actionService.getByJobId(jobId).finishedAt()); - Assertions.assertNull(actionService.getByJobId(jobId).afterSnapshotId()); - - // Late finish listener still records the real terminal snapshot. - actionService.markFinished(jobId, 200L, System.currentTimeMillis()); - Assertions.assertEquals(200L, actionService.getByJobId(jobId).afterSnapshotId()); - Assertions.assertEquals(reclaimFinishedAt, actionService.getByJobId(jobId).finishedAt()); + actionService.markFinished(jobId, tableId, policyId, null, reclaimFinishedAt); + Assertions.assertEquals( + reclaimFinishedAt, actionService.getByJobIdAndTableId(jobId, tableId).finishedAt()); + Assertions.assertNull(actionService.getByJobIdAndTableId(jobId, tableId).afterSnapshotId()); + + actionService.markFinished(jobId, tableId, policyId, 200L, System.currentTimeMillis()); + Assertions.assertEquals( + 200L, actionService.getByJobIdAndTableId(jobId, tableId).afterSnapshotId()); + Assertions.assertEquals( + reclaimFinishedAt, actionService.getByJobIdAndTableId(jobId, tableId).finishedAt()); Assertions.assertThrows( IllegalStateException.class, - () -> actionService.markFinished(jobId, 201L, System.currentTimeMillis())); + () -> + actionService.markFinished(jobId, tableId, policyId, 201L, System.currentTimeMillis())); + } + + @TestTemplate + void testBindJobIdIdempotentWhenAlreadyBoundToSameJob() { + TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); + TableMaintenanceConfig config = TableMaintenanceConfig.forTest(1, 1, 1, 500L, 1_000L, 4, 10); + long tableId = 31L; + long policyId = 41L; + long jobId = 5100L; + OccupancyClaim claim = actions.insertInFlight(tableId, policyId, config); + actions.bindJobId(claim, jobId); + // Spark lazy-bound first; submit-side bind must not fail. + actions.bindJobId(claim, jobId); + Assertions.assertEquals(jobId, actions.getByJobIdAndTableId(jobId, tableId).jobId()); + } + + @TestTemplate + void testBindJobIdFailsWhenBoundToDifferentJob() { + TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); + TableMaintenanceConfig config = TableMaintenanceConfig.forTest(1, 1, 1, 500L, 1_000L, 4, 10); + OccupancyClaim claim = actions.insertInFlight(32L, 42L, config); + actions.bindJobId(claim, 5200L); + Assertions.assertThrows(IllegalStateException.class, () -> actions.bindJobId(claim, 5201L)); + } + + @TestTemplate + void testBindJobIdIfUnboundLazyBindsSoleClaim() { + TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); + TableMaintenanceConfig config = TableMaintenanceConfig.forTest(1, 1, 1, 500L, 1_000L, 4, 10); + long tableId = 33L; + long policyId = 43L; + long jobId = 5300L; + actions.insertInFlight(tableId, policyId, config); + + Assertions.assertTrue(actions.bindJobIdIfUnbound(tableId, policyId, jobId)); + Assertions.assertEquals( + jobId, actions.getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyId).jobId()); + // Already bound to same job — still reports success via idempotent bind. + Assertions.assertTrue(actions.bindJobIdIfUnbound(tableId, policyId, jobId)); + } + + @TestTemplate + void testBindJobIdIfUnboundDisambiguatesPoliciesOnSameTable() { + TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); + TableMaintenanceConfig config = TableMaintenanceConfig.forTest(1, 1, 1, 500L, 1_000L, 4, 10); + long tableId = 34L; + long policyA = 44L; + long policyB = 45L; + long jobId = 5400L; + // Two policies → two soft slots on the same table; policyId selects which claim to bind. + actions.insertInFlight(tableId, policyA, config); + actions.insertInFlight(tableId, policyB, config); + + Assertions.assertTrue(actions.bindJobIdIfUnbound(tableId, policyA, jobId)); + Assertions.assertEquals( + jobId, actions.getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyA).jobId()); + Assertions.assertNull(actions.getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyB)); + Assertions.assertTrue(actions.bindJobIdIfUnbound(tableId, policyB, jobId)); + Assertions.assertEquals( + jobId, actions.getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyB).jobId()); + } + + @TestTemplate + void testOneJobMultiplePoliciesOnSameTableRecordAndFinishPerPolicy() { + TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); + TableMaintenanceConfig config = TableMaintenanceConfig.forTest(1, 1, 1, 500L, 1_000L, 4, 10); + long tableId = 36L; + long policyA = 46L; + long policyB = 47L; + long jobId = 5600L; + OccupancyClaim claimA = actions.insertInFlight(tableId, policyA, config); + OccupancyClaim claimB = actions.insertInFlight(tableId, policyB, config); + actions.bindJobId(claimA, jobId); + actions.bindJobId(claimB, jobId); + + actions.recordBeforeSnapshot(jobId, tableId, policyA, 11L); + actions.recordBeforeSnapshot(jobId, tableId, policyB, 22L); + Assertions.assertEquals( + 11L, actions.getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyA).beforeSnapshotId()); + Assertions.assertEquals( + 22L, actions.getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyB).beforeSnapshotId()); + + long finishedAt = System.currentTimeMillis(); + actions.markFinished(jobId, tableId, policyA, null, finishedAt); + Assertions.assertEquals( + finishedAt, actions.getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyA).finishedAt()); + Assertions.assertEquals( + 0L, actions.getByJobIdAndTableIdAndPolicyId(jobId, tableId, policyB).finishedAt()); + } + + @TestTemplate + void testBindJobIdIfUnboundReturnsFalseWhenNoClaim() { + TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); + Assertions.assertFalse(actions.bindJobIdIfUnbound(35L, 55L, 5500L)); } @TestTemplate diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/web/rest/TestTableMaintenanceActionOperations.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/web/rest/TestTableMaintenanceActionOperations.java new file mode 100644 index 0000000000..ad86568cd1 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/web/rest/TestTableMaintenanceActionOperations.java @@ -0,0 +1,234 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.maintenance.web.rest; + +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.ArgumentMatchers.isNull; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.function.LongFunction; +import javax.servlet.http.HttpServletRequest; +import javax.ws.rs.core.Response; +import org.apache.gravitino.maintenance.MaintenanceMetricsService; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionFinishRequest; +import org.apache.gravitino.maintenance.dto.requests.MaintenanceActionMetricsRequest; +import org.apache.gravitino.maintenance.storage.po.TableMaintenanceActionPO; +import org.apache.gravitino.maintenance.storage.service.TableMaintenanceActionMetaService; +import org.apache.gravitino.maintenance.storage.service.TableSnapshotMetricsMetaService; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +/** Unit tests for {@link TableMaintenanceActionOperations}. */ +class TestTableMaintenanceActionOperations { + + private static final String METALAKE = "ml"; + private static final long JOB_ID = 42L; + private static final long TABLE_ID = 10L; + private static final long POLICY_ID = 20L; + + private TableMaintenanceActionMetaService actions; + private TableSnapshotMetricsMetaService metrics; + private TableMaintenanceActionOperations ops; + + @BeforeEach + void setUp() { + actions = mock(TableMaintenanceActionMetaService.class); + metrics = mock(TableSnapshotMetricsMetaService.class); + ops = operations(tableId -> METALAKE); + ops.setHttpRequest(mock(HttpServletRequest.class)); + } + + @Test + void testReportBeforePersistsAndRecordsBeforeSnapshot() { + when(actions.getByJobIdAndTableIdAndPolicyId(JOB_ID, TABLE_ID, POLICY_ID)) + .thenReturn(action(TABLE_ID, POLICY_ID)); + + Response response = + ops.reportBeforeMetrics( + METALAKE, + JOB_ID, + /*finish=*/ false, + new MaintenanceActionMetricsRequest(TABLE_ID, POLICY_ID, 100L, "{\"m\":1}")); + + Assertions.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + verify(metrics).insertMetrics(TABLE_ID, POLICY_ID, 100L, "{\"m\":1}"); + verify(actions).recordBeforeSnapshot(JOB_ID, TABLE_ID, POLICY_ID, 100L); + verify(actions, never()).markFinished(anyLong(), anyLong(), anyLong(), anyLong(), anyLong()); + } + + @Test + void testReportBeforeWithFinishClosesSoftSlotWithoutAfterSnapshot() { + when(actions.getByJobIdAndTableIdAndPolicyId(JOB_ID, TABLE_ID, POLICY_ID)) + .thenReturn(action(TABLE_ID, POLICY_ID)); + + Response response = + ops.reportBeforeMetrics( + METALAKE, + JOB_ID, + /*finish=*/ true, + new MaintenanceActionMetricsRequest(TABLE_ID, POLICY_ID, 100L, "{\"m\":1}")); + + Assertions.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + verify(actions).markFinished(eq(JOB_ID), eq(TABLE_ID), eq(POLICY_ID), isNull(), anyLong()); + } + + @Test + void testReportAfterWithFinishClosesSoftSlotWithAfterSnapshot() { + when(actions.getByJobIdAndTableIdAndPolicyId(JOB_ID, TABLE_ID, POLICY_ID)) + .thenReturn(action(TABLE_ID, POLICY_ID)); + + Response response = + ops.reportAfterMetrics( + METALAKE, + JOB_ID, + /*finish=*/ true, + new MaintenanceActionMetricsRequest(TABLE_ID, POLICY_ID, 200L, "{\"m\":2}")); + + Assertions.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + verify(metrics).insertMetrics(TABLE_ID, POLICY_ID, 200L, "{\"m\":2}"); + verify(actions).markFinished(eq(JOB_ID), eq(TABLE_ID), eq(POLICY_ID), eq(200L), anyLong()); + } + + @Test + void testReportBeforeMissingActionReturns404() { + when(actions.getByJobIdAndTableIdAndPolicyId(JOB_ID, TABLE_ID, POLICY_ID)).thenReturn(null); + when(actions.bindJobIdIfUnbound(TABLE_ID, POLICY_ID, JOB_ID)).thenReturn(false); + + Response response = + ops.reportBeforeMetrics( + METALAKE, + JOB_ID, + /*finish=*/ false, + new MaintenanceActionMetricsRequest(TABLE_ID, POLICY_ID, 100L, "{\"m\":1}")); + + Assertions.assertEquals(Response.Status.NOT_FOUND.getStatusCode(), response.getStatus()); + verify(metrics, never()).insertMetrics(anyLong(), anyLong(), anyLong(), anyString()); + } + + @Test + void testReportBeforeLazyBindsUnboundClaim() { + when(actions.getByJobIdAndTableIdAndPolicyId(JOB_ID, TABLE_ID, POLICY_ID)) + .thenReturn(null) + .thenReturn(action(TABLE_ID, POLICY_ID)); + when(actions.bindJobIdIfUnbound(TABLE_ID, POLICY_ID, JOB_ID)).thenReturn(true); + + Response response = + ops.reportBeforeMetrics( + METALAKE, + JOB_ID, + /*finish=*/ false, + new MaintenanceActionMetricsRequest(TABLE_ID, POLICY_ID, 100L, "{\"m\":1}")); + + Assertions.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + verify(actions).bindJobIdIfUnbound(TABLE_ID, POLICY_ID, JOB_ID); + verify(metrics).insertMetrics(TABLE_ID, POLICY_ID, 100L, "{\"m\":1}"); + verify(actions).recordBeforeSnapshot(JOB_ID, TABLE_ID, POLICY_ID, 100L); + } + + @Test + void testFinishLazyBindsUnboundClaim() { + when(actions.getByJobIdAndTableIdAndPolicyId(JOB_ID, TABLE_ID, POLICY_ID)) + .thenReturn(null) + .thenReturn(action(TABLE_ID, POLICY_ID)); + when(actions.bindJobIdIfUnbound(TABLE_ID, POLICY_ID, JOB_ID)).thenReturn(true); + + Response response = + ops.finishAction( + METALAKE, JOB_ID, new MaintenanceActionFinishRequest(TABLE_ID, POLICY_ID, null)); + + Assertions.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + verify(actions).bindJobIdIfUnbound(TABLE_ID, POLICY_ID, JOB_ID); + verify(actions).markFinished(eq(JOB_ID), eq(TABLE_ID), eq(POLICY_ID), isNull(), anyLong()); + } + + @Test + void testReportAfterAlreadyFinishedReturns404() { + when(actions.getByJobIdAndTableIdAndPolicyId(JOB_ID, TABLE_ID, POLICY_ID)) + .thenReturn(action(TABLE_ID, POLICY_ID, /*finishedAt=*/ 9L)); + + Response response = + ops.reportAfterMetrics( + METALAKE, + JOB_ID, + /*finish=*/ false, + new MaintenanceActionMetricsRequest(TABLE_ID, POLICY_ID, 100L, "{\"m\":1}")); + + Assertions.assertEquals(Response.Status.NOT_FOUND.getStatusCode(), response.getStatus()); + verify(metrics, never()).insertMetrics(anyLong(), anyLong(), anyLong(), anyString()); + } + + @Test + void testReportBeforeNullBodyReturns400() { + Response response = ops.reportBeforeMetrics(METALAKE, JOB_ID, /*finish=*/ false, null); + Assertions.assertEquals(Response.Status.BAD_REQUEST.getStatusCode(), response.getStatus()); + } + + @Test + void testFinishClosesSoftSlotWithoutMetrics() { + when(actions.getByJobIdAndTableIdAndPolicyId(JOB_ID, TABLE_ID, POLICY_ID)) + .thenReturn(action(TABLE_ID, POLICY_ID)); + + Response response = + ops.finishAction( + METALAKE, JOB_ID, new MaintenanceActionFinishRequest(TABLE_ID, POLICY_ID, null)); + + Assertions.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + verify(metrics, never()).insertMetrics(anyLong(), anyLong(), anyLong(), anyString()); + verify(actions).markFinished(eq(JOB_ID), eq(TABLE_ID), eq(POLICY_ID), isNull(), anyLong()); + } + + @Test + void testFinishWithAfterSnapshotId() { + when(actions.getByJobIdAndTableIdAndPolicyId(JOB_ID, TABLE_ID, POLICY_ID)) + .thenReturn(action(TABLE_ID, POLICY_ID)); + + Response response = + ops.finishAction( + METALAKE, JOB_ID, new MaintenanceActionFinishRequest(TABLE_ID, POLICY_ID, 300L)); + + Assertions.assertEquals(Response.Status.OK.getStatusCode(), response.getStatus()); + verify(actions).markFinished(eq(JOB_ID), eq(TABLE_ID), eq(POLICY_ID), eq(300L), anyLong()); + } + + private TableMaintenanceActionOperations operations(LongFunction<String> lookup) { + return new TableMaintenanceActionOperations( + new MaintenanceMetricsService(actions, metrics, lookup)); + } + + private static TableMaintenanceActionPO action(long tableId, long policyId) { + return action(tableId, policyId, 0L); + } + + private static TableMaintenanceActionPO action(long tableId, long policyId, long finishedAt) { + return TableMaintenanceActionPO.builder() + .withId(1L) + .withJobId(JOB_ID) + .withTableId(tableId) + .withPolicyId(policyId) + .withFinishedAt(finishedAt) + .withCreatedAt(1L) + .build(); + } +} diff --git a/scripts/h2/schema-2.0.0-h2.sql b/scripts/h2/schema-2.0.0-h2.sql index 901f3c99cd..9343462fe6 100644 --- a/scripts/h2/schema-2.0.0-h2.sql +++ b/scripts/h2/schema-2.0.0-h2.sql @@ -743,7 +743,7 @@ CREATE TABLE IF NOT EXISTS `table_maintenance_action` ( `created_at` BIGINT NOT NULL COMMENT 'creation timestamp in milliseconds', PRIMARY KEY (`id`) ) COMMENT='table maintenance action metadata'; -CREATE UNIQUE INDEX IF NOT EXISTS `uk_tma_job_id` ON `table_maintenance_action`(`job_id`); +CREATE INDEX IF NOT EXISTS `idx_tma_job_id` ON `table_maintenance_action`(`job_id`); -- Soft slot: finished_at=0 is exclusive per table/policy (one in-flight finish). CREATE UNIQUE INDEX IF NOT EXISTS `uk_tma_table_policy_finished` ON `table_maintenance_action`(`table_id`, `policy_id`, `finished_at`); CREATE INDEX IF NOT EXISTS `idx_tma_before_sid` ON `table_maintenance_action`(`table_id`, `before_snapshot_id`); diff --git a/scripts/h2/upgrade-1.3.0-to-2.0.0-h2.sql b/scripts/h2/upgrade-1.3.0-to-2.0.0-h2.sql index ffc5b9ca11..6d84961017 100644 --- a/scripts/h2/upgrade-1.3.0-to-2.0.0-h2.sql +++ b/scripts/h2/upgrade-1.3.0-to-2.0.0-h2.sql @@ -164,7 +164,7 @@ CREATE TABLE IF NOT EXISTS `table_maintenance_action` ( `created_at` BIGINT NOT NULL COMMENT 'creation timestamp in milliseconds', PRIMARY KEY (`id`) ) COMMENT='table maintenance action metadata'; -CREATE UNIQUE INDEX IF NOT EXISTS `uk_tma_job_id` ON `table_maintenance_action`(`job_id`); +CREATE INDEX IF NOT EXISTS `idx_tma_job_id` ON `table_maintenance_action`(`job_id`); -- Soft slot: finished_at=0 is exclusive per table/policy (one in-flight finish). CREATE UNIQUE INDEX IF NOT EXISTS `uk_tma_table_policy_finished` ON `table_maintenance_action`(`table_id`, `policy_id`, `finished_at`); CREATE INDEX IF NOT EXISTS `idx_tma_before_sid` ON `table_maintenance_action`(`table_id`, `before_snapshot_id`); diff --git a/scripts/mysql/schema-2.0.0-mysql.sql b/scripts/mysql/schema-2.0.0-mysql.sql index db39e065ba..d8110fa483 100644 --- a/scripts/mysql/schema-2.0.0-mysql.sql +++ b/scripts/mysql/schema-2.0.0-mysql.sql @@ -729,7 +729,7 @@ CREATE TABLE IF NOT EXISTS `table_maintenance_action` ( `finished_at` BIGINT(20) UNSIGNED NOT NULL DEFAULT 0 COMMENT 'finish millis, 0 = soft in-flight slot (one per table/policy)', `created_at` BIGINT(20) UNSIGNED NOT NULL COMMENT 'creation timestamp in milliseconds', PRIMARY KEY (`id`), - UNIQUE KEY `uk_tma_job_id` (`job_id`), + KEY `idx_tma_job_id` (`job_id`), -- Soft slot: finished_at=0 is exclusive per table/policy (one in-flight finish). UNIQUE KEY `uk_tma_table_policy_finished` (`table_id`, `policy_id`, `finished_at`), KEY `idx_tma_before_sid` (`table_id`, `before_snapshot_id`), diff --git a/scripts/mysql/upgrade-1.3.0-to-2.0.0-mysql.sql b/scripts/mysql/upgrade-1.3.0-to-2.0.0-mysql.sql index 26677f167e..07c35e345c 100644 --- a/scripts/mysql/upgrade-1.3.0-to-2.0.0-mysql.sql +++ b/scripts/mysql/upgrade-1.3.0-to-2.0.0-mysql.sql @@ -610,7 +610,7 @@ CREATE TABLE IF NOT EXISTS `table_maintenance_action` ( `finished_at` BIGINT(20) UNSIGNED NOT NULL DEFAULT 0 COMMENT 'finish millis, 0 = soft in-flight slot (one per table/policy)', `created_at` BIGINT(20) UNSIGNED NOT NULL COMMENT 'creation timestamp in milliseconds', PRIMARY KEY (`id`), - UNIQUE KEY `uk_tma_job_id` (`job_id`), + KEY `idx_tma_job_id` (`job_id`), -- Soft slot: finished_at=0 is exclusive per table/policy (one in-flight finish). UNIQUE KEY `uk_tma_table_policy_finished` (`table_id`, `policy_id`, `finished_at`), KEY `idx_tma_before_sid` (`table_id`, `before_snapshot_id`), diff --git a/scripts/postgresql/schema-2.0.0-postgresql.sql b/scripts/postgresql/schema-2.0.0-postgresql.sql index 19868d56db..81b2e8cc71 100644 --- a/scripts/postgresql/schema-2.0.0-postgresql.sql +++ b/scripts/postgresql/schema-2.0.0-postgresql.sql @@ -1273,7 +1273,7 @@ CREATE TABLE IF NOT EXISTS table_maintenance_action ( finished_at BIGINT NOT NULL DEFAULT 0, created_at BIGINT NOT NULL ); -CREATE UNIQUE INDEX IF NOT EXISTS uk_tma_job_id ON table_maintenance_action (job_id); +CREATE INDEX IF NOT EXISTS idx_tma_job_id ON table_maintenance_action (job_id); -- Soft slot: finished_at=0 is exclusive per table/policy (one in-flight finish). CREATE UNIQUE INDEX IF NOT EXISTS uk_tma_table_policy_finished ON table_maintenance_action (table_id, policy_id, finished_at); CREATE INDEX IF NOT EXISTS idx_tma_before_sid ON table_maintenance_action (table_id, before_snapshot_id); diff --git a/scripts/postgresql/upgrade-1.3.0-to-2.0.0-postgresql.sql b/scripts/postgresql/upgrade-1.3.0-to-2.0.0-postgresql.sql index b6deb5b273..0cf799dd09 100644 --- a/scripts/postgresql/upgrade-1.3.0-to-2.0.0-postgresql.sql +++ b/scripts/postgresql/upgrade-1.3.0-to-2.0.0-postgresql.sql @@ -248,7 +248,7 @@ CREATE TABLE IF NOT EXISTS table_maintenance_action ( finished_at BIGINT NOT NULL DEFAULT 0, created_at BIGINT NOT NULL ); -CREATE UNIQUE INDEX IF NOT EXISTS uk_tma_job_id ON table_maintenance_action (job_id); +CREATE INDEX IF NOT EXISTS idx_tma_job_id ON table_maintenance_action (job_id); -- Soft slot: finished_at=0 is exclusive per table/policy (one in-flight finish). CREATE UNIQUE INDEX IF NOT EXISTS uk_tma_table_policy_finished ON table_maintenance_action (table_id, policy_id, finished_at); CREATE INDEX IF NOT EXISTS idx_tma_before_sid ON table_maintenance_action (table_id, before_snapshot_id); diff --git a/settings.gradle.kts b/settings.gradle.kts index 4d15b84397..0076ce9d82 100644 --- a/settings.gradle.kts +++ b/settings.gradle.kts @@ -136,6 +136,7 @@ include(":plugins:idp-basic") include(":plugins:k8s-job-executor") include(":plugins:tms") include( + ":maintenance:maintenance-common", ":maintenance:optimizer-api", ":maintenance:updaters", ":maintenance:optimizer",
