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
The following commit(s) were added to refs/heads/main by this push:
new 189ba4dfcd [#13777] feat(tms): In-job S0 gates, OccupancyGate, and
REST Feature lifecycle (#13841)
189ba4dfcd is described below
commit 189ba4dfcd53936893c3f6b7972eae79b26cdca6
Author: MaSai <[email protected]>
AuthorDate: Thu Oct 15 10:10:09 2026 +0800
[#13777] feat(tms): In-job S0 gates, OccupancyGate, and REST Feature
lifecycle (#13841)
### What changes were proposed in this pull request?
Combines in-job gates and OccupancyGate / REST Feature work on top of
merged [#13774](https://github.com/apache/gravitino/pull/13774)
(formerly tracked as #13816, folded into #13777):
**In-job gates**
- Package `maintenance.injob`: S0 sample → snapshot-drift → recommend
gates
- `TableMaintenanceInJobSampler` / `TableMaintenanceInJobPipeline`
(`gateCronJob` / `gateCommitChainType`)
- `InJobRecommender` + `RecommendDecision` (injected in tests; no
production placeholder recommender)
- Persist sampled metrics via idempotent `insertMetrics`
(first-writer-wins on unique race) and set `before_snapshot_id` (same-S0
retry safe)
- Unit tests bind soft-slot rows with `OccupancyClaim` (`insertInFlight`
+ `bindJobId`)
**OccupancyGate + REST Feature**
- `OccupancyGate` over soft-slot occupancy: unbound TTL evaluated on the
**database** clock (`hasLiveUnboundClaim`); bound rows consult
`JobMetaService` with no `finished_at` backfill
- `TableMaintenanceRESTFeature` starts the three db-scheduler pools when
the TMS Jersey extension package is enabled, and stops them on Jersey
`DESTROY_FINISHED` (before entity-store `DataSource` close);
configure-after-stop re-registers the destroy listener
- Jersey / `server-common` dependencies for the Feature SPI
- Unit tests for `OccupancyGate` (including DB-clock unbound expiry) and
Feature destroy listener
**Out of scope:** Four maintenance type executors, OCC commit-chain,
PolicyExpand / Spark submit / IRC listener pipelines.
### Why are the changes needed?
In-job S0/drift/recommend gates and control-plane occupancy / Feature
lifecycle for Epic #13760.
Fix: #13777
### Does this PR introduce _any_ user-facing change?
No public REST resources yet. Enabling the TMS extension package starts
in-process scheduler pools via `TableMaintenanceRESTFeature`.
### How was this patch tested?
```bash
./gradlew :plugins:tms:test --tests
'org.apache.gravitino.maintenance.injob.*' \
--tests 'org.apache.gravitino.maintenance.occupancy.*' \
--tests 'org.apache.gravitino.maintenance.web.rest.feature.*' \
--tests 'org.apache.gravitino.maintenance.storage.service.*' -PskipITs
```
---------
Co-authored-by: Cursor <[email protected]>
---
plugins/tms/build.gradle.kts | 4 +
.../maintenance/injob/InJobGateOutcome.java | 99 +++++++++
.../maintenance/injob/InJobRecommender.java | 39 ++++
.../maintenance/injob/RecommendDecision.java | 33 +++
.../gravitino/maintenance/injob/SnapshotHead.java | 33 +++
.../injob/TableMaintenanceInJobPipeline.java | 105 ++++++++++
.../injob/TableMaintenanceInJobSampler.java | 105 ++++++++++
.../maintenance/occupancy/OccupancyGate.java | 183 +++++++++++++++++
.../TableMaintenanceActionBaseSQLProvider.java | 3 +-
.../service/TableMaintenanceActionMetaService.java | 43 +++-
.../service/TableSnapshotMetricsMetaService.java | 32 ++-
.../rest/feature/TableMaintenanceRESTFeature.java | 130 ++++++++++++
.../injob/TestTableMaintenanceInJobPipeline.java | 99 +++++++++
.../injob/TestTableMaintenanceInJobSampler.java | 107 ++++++++++
.../maintenance/occupancy/TestOccupancyGate.java | 221 +++++++++++++++++++++
.../TestTableMaintenanceActionMetaService.java | 3 +
.../feature/TestTableMaintenanceRESTFeature.java | 40 ++++
17 files changed, 1272 insertions(+), 7 deletions(-)
diff --git a/plugins/tms/build.gradle.kts b/plugins/tms/build.gradle.kts
index 835eb0fea6..b12be69734 100644
--- a/plugins/tms/build.gradle.kts
+++ b/plugins/tms/build.gradle.kts
@@ -29,10 +29,13 @@ dependencies {
implementation(project(":api"))
implementation(project(":common"))
implementation(project(":core"))
+ implementation(project(":server-common"))
+ implementation(libs.bundles.jersey)
implementation(libs.db.scheduler)
implementation(libs.guava)
implementation(libs.mybatis)
+ implementation(libs.servlet)
compileOnly(libs.lombok)
compileOnly(libs.slf4j.api)
@@ -41,6 +44,7 @@ dependencies {
testImplementation(project(":core", "testArtifacts"))
testImplementation(project(":integration-test-common", "testArtifacts"))
testImplementation(project(":server"))
+ testImplementation(project(":server-common"))
testImplementation(libs.junit.jupiter.api)
testImplementation(libs.junit.jupiter.params)
testImplementation(libs.mockito.core)
diff --git
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/InJobGateOutcome.java
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/InJobGateOutcome.java
new file mode 100644
index 0000000000..61cb521881
--- /dev/null
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/InJobGateOutcome.java
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.injob;
+
+import com.google.common.base.Preconditions;
+import javax.annotation.Nullable;
+
+/** Result of the in-job sample / drift / recommend gates (§5.2.2). */
+public final class InJobGateOutcome {
+
+ /** Why maintenance should not continue. */
+ public enum StopReason {
+ /** HEAD snapshot id differs from {@code S0}; do not re-sample in a loop.
*/
+ SNAPSHOT_DRIFT,
+ /** Recommender decided not to maintain; end early without running work. */
+ RECOMMENDER_SKIP
+ }
+
+ private final boolean continueMaintenance;
+ @Nullable private final StopReason stopReason;
+ private final long sampledSnapshotId;
+ private final long headSnapshotId;
+
+ private InJobGateOutcome(
+ boolean continueMaintenance,
+ @Nullable StopReason stopReason,
+ long sampledSnapshotId,
+ long headSnapshotId) {
+ this.continueMaintenance = continueMaintenance;
+ this.stopReason = stopReason;
+ this.sampledSnapshotId = sampledSnapshotId;
+ this.headSnapshotId = headSnapshotId;
+ }
+
+ /**
+ * @param sampledSnapshotId {@code S0}
+ * @return outcome when maintenance may proceed
+ */
+ public static InJobGateOutcome continueMaintenance(long sampledSnapshotId) {
+ return new InJobGateOutcome(true, null, sampledSnapshotId,
sampledSnapshotId);
+ }
+
+ /**
+ * @param reason why the job or type should stop
+ * @param sampledSnapshotId {@code S0}
+ * @param headSnapshotId current HEAD when stopping
+ * @return outcome when maintenance must not continue
+ */
+ public static InJobGateOutcome stop(
+ StopReason reason, long sampledSnapshotId, long headSnapshotId) {
+ Preconditions.checkNotNull(reason);
+ return new InJobGateOutcome(false, reason, sampledSnapshotId,
headSnapshotId);
+ }
+
+ /**
+ * @return whether the Spark Job should run work for this type
+ */
+ public boolean continueMaintenance() {
+ return continueMaintenance;
+ }
+
+ /**
+ * @return stop reason when {@link #continueMaintenance()} is false
+ */
+ @Nullable
+ public StopReason stopReason() {
+ return stopReason;
+ }
+
+ /**
+ * @return sampled snapshot id ({@code S0})
+ */
+ public long sampledSnapshotId() {
+ return sampledSnapshotId;
+ }
+
+ /**
+ * @return HEAD snapshot id observed for the gate
+ */
+ public long headSnapshotId() {
+ return headSnapshotId;
+ }
+}
diff --git
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/InJobRecommender.java
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/InJobRecommender.java
new file mode 100644
index 0000000000..dd05f934e7
--- /dev/null
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/InJobRecommender.java
@@ -0,0 +1,39 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.injob;
+
+/**
+ * Decides whether maintenance should run after metrics were sampled at {@code
S0}.
+ *
+ * <p>Production implementations come with the four maintenance executors;
callers inject the
+ * strategy for the current type.
+ */
+@FunctionalInterface
+public interface InJobRecommender {
+
+ /**
+ * @param tableId table entity id
+ * @param policyId policy entity id
+ * @param sampledSnapshotId Iceberg snapshot id at sample time
+ * @param metricsJson serialized metrics for that snapshot
+ * @return maintain or skip
+ */
+ RecommendDecision recommend(
+ long tableId, long policyId, long sampledSnapshotId, String metricsJson);
+}
diff --git
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/RecommendDecision.java
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/RecommendDecision.java
new file mode 100644
index 0000000000..eff0d58650
--- /dev/null
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/RecommendDecision.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.injob;
+
+/**
+ * Whether the in-job path should run maintenance work after a successful
{@code S0} sample
+ * (§5.2.2).
+ *
+ * <p>Type-specific strategy recommenders are wired in a follow-up subtask;
this enum only models
+ * the gate outcome.
+ */
+public enum RecommendDecision {
+ /** Proceed to execute maintenance work. */
+ MAINTAIN,
+ /** End this job or type early; do not run maintenance work. */
+ SKIP
+}
diff --git
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/SnapshotHead.java
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/SnapshotHead.java
new file mode 100644
index 0000000000..04e6f089b8
--- /dev/null
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/SnapshotHead.java
@@ -0,0 +1,33 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.injob;
+
+/**
+ * Supplies the Iceberg table HEAD snapshot id from inside a Spark Job (§4.3,
§5.2.2).
+ *
+ * <p>Implementations typically delegate to {@code Table#currentSnapshot()}.
+ */
+@FunctionalInterface
+public interface SnapshotHead {
+
+ /**
+ * @return current Iceberg snapshot id at table HEAD
+ */
+ long currentSnapshotId();
+}
diff --git
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/TableMaintenanceInJobPipeline.java
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/TableMaintenanceInJobPipeline.java
new file mode 100644
index 0000000000..9cef15ca4e
--- /dev/null
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/TableMaintenanceInJobPipeline.java
@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.injob;
+
+import com.google.common.base.Preconditions;
+
+/**
+ * Cron / commit-chain in-job gates: sample at {@code S0}, stop on snapshot
drift, then apply the
+ * recommend maintain? gate (§5.2.2). Does not execute the four maintenance
types.
+ */
+public class TableMaintenanceInJobPipeline {
+
+ private final TableMaintenanceInJobSampler sampler;
+ private final InJobRecommender recommender;
+
+ /**
+ * @param sampler S0 sample + drift gate
+ * @param recommender maintain vs skip after a clean sample
+ */
+ public TableMaintenanceInJobPipeline(
+ TableMaintenanceInJobSampler sampler, InJobRecommender recommender) {
+ this.sampler = Preconditions.checkNotNull(sampler);
+ this.recommender = Preconditions.checkNotNull(recommender);
+ }
+
+ /**
+ * Runs sample → drift → recommend for a cron Spark Job.
+ *
+ * @param jobId Gravitino {@code job_run_meta.job_run_id}
+ * @param tableId table entity id
+ * @param policyId policy entity id
+ * @param sampledSnapshotId {@code S0}
+ * @param metricsJson metrics payload
+ * @param head current HEAD supplier after persistence
+ * @return continue only when HEAD matches {@code S0} and recommender says
maintain
+ */
+ public InJobGateOutcome gateCronJob(
+ long jobId,
+ long tableId,
+ long policyId,
+ long sampledSnapshotId,
+ String metricsJson,
+ SnapshotHead head) {
+ InJobGateOutcome afterSample =
+ sampler.sampleAndGateCronJob(
+ jobId, tableId, policyId, sampledSnapshotId, metricsJson, head);
+ if (!afterSample.continueMaintenance()) {
+ return afterSample;
+ }
+
+ RecommendDecision decision =
+ recommender.recommend(tableId, policyId, sampledSnapshotId,
metricsJson);
+ if (decision == RecommendDecision.SKIP) {
+ return InJobGateOutcome.stop(
+ InJobGateOutcome.StopReason.RECOMMENDER_SKIP,
+ sampledSnapshotId,
+ afterSample.headSnapshotId());
+ }
+ return afterSample;
+ }
+
+ /**
+ * Drift then recommend for one type inside a commit-chain job (§5.2.3).
Caller ends only the
+ * current type on stop; does not re-sample.
+ *
+ * @param tableId table entity id
+ * @param policyId policy entity id
+ * @param sampledSnapshotId {@code S0} for this type
+ * @param metricsJson metrics for this type's sample
+ * @param head current HEAD after this type's sample persistence
+ * @return continue vs stop for this type
+ */
+ public InJobGateOutcome gateCommitChainType(
+ long tableId, long policyId, long sampledSnapshotId, String metricsJson,
SnapshotHead head) {
+ InJobGateOutcome afterDrift =
sampler.gateCommitChainType(sampledSnapshotId, head);
+ if (!afterDrift.continueMaintenance()) {
+ return afterDrift;
+ }
+ RecommendDecision decision =
+ recommender.recommend(tableId, policyId, sampledSnapshotId,
metricsJson);
+ if (decision == RecommendDecision.SKIP) {
+ return InJobGateOutcome.stop(
+ InJobGateOutcome.StopReason.RECOMMENDER_SKIP,
+ sampledSnapshotId,
+ afterDrift.headSnapshotId());
+ }
+ return afterDrift;
+ }
+}
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
new file mode 100644
index 0000000000..9c96638434
--- /dev/null
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/injob/TableMaintenanceInJobSampler.java
@@ -0,0 +1,105 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.injob;
+
+import com.google.common.base.Preconditions;
+import
org.apache.gravitino.maintenance.storage.service.TableMaintenanceActionMetaService;
+import
org.apache.gravitino.maintenance.storage.service.TableSnapshotMetricsMetaService;
+
+/**
+ * In-Spark sampling path for TMS cron jobs (§5.2.2): sample once at {@code
S0}, persist metrics,
+ * record {@code before_snapshot_id}, then gate on snapshot drift without
re-sampling.
+ */
+public class TableMaintenanceInJobSampler {
+
+ private final TableSnapshotMetricsMetaService metricsMetaService;
+ private final TableMaintenanceActionMetaService actionMetaService;
+
+ /** Creates a sampler using the default relational meta services. */
+ public TableMaintenanceInJobSampler() {
+ this(
+ TableSnapshotMetricsMetaService.getInstance(),
+ TableMaintenanceActionMetaService.getInstance());
+ }
+
+ /**
+ * @param metricsMetaService metrics DAO
+ * @param actionMetaService maintenance action DAO
+ */
+ public TableMaintenanceInJobSampler(
+ TableSnapshotMetricsMetaService metricsMetaService,
+ TableMaintenanceActionMetaService actionMetaService) {
+ this.metricsMetaService = Preconditions.checkNotNull(metricsMetaService);
+ this.actionMetaService = Preconditions.checkNotNull(actionMetaService);
+ }
+
+ /**
+ * Persists metrics for {@code S0}, sets {@code before_snapshot_id} on the
maintenance action row,
+ * and returns whether the cron Spark Job should continue.
+ *
+ * <p>If HEAD has drifted after the sample, the caller must end the job
early and must not
+ * re-sample in a loop (§4.3).
+ *
+ * @param jobId Gravitino {@code job_run_meta.job_run_id} for the in-flight
maintenance row
+ * @param tableId table entity id
+ * @param policyId policy entity id
+ * @param sampledSnapshotId Iceberg snapshot id at sample time ({@code S0})
+ * @param metricsJson serialized metrics payload for {@code
table_snapshot_metrics}
+ * @param head supplies current HEAD after persistence
+ * @return continue vs drift stop outcome
+ */
+ public InJobGateOutcome sampleAndGateCronJob(
+ long jobId,
+ long tableId,
+ long policyId,
+ long sampledSnapshotId,
+ String metricsJson,
+ SnapshotHead head) {
+ Preconditions.checkNotNull(metricsJson, "metricsJson");
+ Preconditions.checkNotNull(head, "head");
+
+ metricsMetaService.insertMetrics(tableId, policyId, sampledSnapshotId,
metricsJson);
+ actionMetaService.recordBeforeSnapshot(jobId, sampledSnapshotId);
+
+ long headSnapshotId = head.currentSnapshotId();
+ if (sampledSnapshotId != headSnapshotId) {
+ return InJobGateOutcome.stop(
+ InJobGateOutcome.StopReason.SNAPSHOT_DRIFT, sampledSnapshotId,
headSnapshotId);
+ }
+ return InJobGateOutcome.continueMaintenance(sampledSnapshotId);
+ }
+
+ /**
+ * Evaluates drift for a commit-chain maintenance type after that type's
sample at {@code S0}
+ * (§5.2.3). When drift is detected the caller ends only the current type
and continues the chain.
+ *
+ * @param sampledSnapshotId {@code S0} for the current type
+ * @param head supplies current HEAD after persistence for this type
+ * @return continue vs drift stop outcome
+ */
+ public InJobGateOutcome gateCommitChainType(long sampledSnapshotId,
SnapshotHead head) {
+ Preconditions.checkNotNull(head, "head");
+ long headSnapshotId = head.currentSnapshotId();
+ if (sampledSnapshotId != headSnapshotId) {
+ return InJobGateOutcome.stop(
+ InJobGateOutcome.StopReason.SNAPSHOT_DRIFT, sampledSnapshotId,
headSnapshotId);
+ }
+ return InJobGateOutcome.continueMaintenance(sampledSnapshotId);
+ }
+}
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
new file mode 100644
index 0000000000..156d75047f
--- /dev/null
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/occupancy/OccupancyGate.java
@@ -0,0 +1,183 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.occupancy;
+
+import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.EnumSet;
+import java.util.List;
+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.mapper.MetalakeMetaMapper;
+import org.apache.gravitino.storage.relational.mapper.TableMetaMapper;
+import org.apache.gravitino.storage.relational.po.MetalakePO;
+import org.apache.gravitino.storage.relational.po.TablePO;
+import org.apache.gravitino.storage.relational.service.JobMetaService;
+import org.apache.gravitino.storage.relational.utils.SessionUtils;
+import org.apache.gravitino.utils.NameIdentifierUtil;
+
+/**
+ * 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} and do <strong>not</strong> backfill
{@code finished_at}.
+ */
+public final class OccupancyGate {
+
+ private static final Set<JobHandle.Status> ACTIVE =
+ EnumSet.of(JobHandle.Status.QUEUED, JobHandle.Status.STARTED,
JobHandle.Status.CANCELLING);
+
+ private final TableMaintenanceActionMetaService actionMetaService;
+ private final long claimTimeoutMs;
+
+ /**
+ * @param actionMetaService maintenance action DAO
+ * @param claimTimeoutMs unbound claim TTL
+ */
+ public OccupancyGate(TableMaintenanceActionMetaService actionMetaService,
long claimTimeoutMs) {
+ Preconditions.checkArgument(claimTimeoutMs > 0, "claimTimeoutMs must be
positive");
+ this.actionMetaService = Preconditions.checkNotNull(actionMetaService);
+ this.claimTimeoutMs = claimTimeoutMs;
+ }
+
+ /**
+ * @param actionMetaService maintenance action DAO
+ */
+ public OccupancyGate(TableMaintenanceActionMetaService actionMetaService) {
+ this(actionMetaService, TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS);
+ }
+
+ /** Uses {@link TableMaintenanceActionMetaService#getInstance()} and default
claim timeout. */
+ public OccupancyGate() {
+ this(TableMaintenanceActionMetaService.getInstance());
+ }
+
+ /**
+ * @param tableId table entity id
+ * @param policyId policy entity id
+ * @return true when any open row is a live unbound claim or an active
Gravitino job
+ */
+ public boolean isInFlight(long tableId, long policyId) {
+ if (actionMetaService.hasLiveUnboundClaim(tableId, policyId,
claimTimeoutMs)) {
+ return true;
+ }
+ for (TableMaintenanceActionPO open :
actionMetaService.listInFlight(tableId, policyId)) {
+ if (open.jobId() == null) {
+ continue;
+ }
+ JobEntity job = loadJob(open.jobId(), tableId);
+ if (job != null && ACTIVE.contains(job.status())) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ /**
+ * @param tableId table entity id
+ * @param policyId policy entity id
+ * @param minIntervalMs resolved minimum interval
+ * @param nowMs current epoch millis used only for cooldown (not unbound TTL)
+ * @return true when the last known finish is still inside the cooldown
window
+ */
+ public boolean withinMinInterval(long tableId, long policyId, long
minIntervalMs, long nowMs) {
+ if (minIntervalMs <= 0) {
+ return false;
+ }
+ Long lastFinished = lastFinishedAt(tableId, policyId);
+ return lastFinished != null && nowMs - lastFinished < minIntervalMs;
+ }
+
+ /**
+ * Expand / submit should skip when occupied or still cooling down.
+ *
+ * @param tableId table entity id
+ * @param policyId policy entity id
+ * @param minIntervalMs resolved minimum interval
+ * @param nowMs current epoch millis used only for cooldown (not unbound TTL)
+ * @return true when enqueue / submit must be skipped
+ */
+ public boolean shouldSkip(long tableId, long policyId, long minIntervalMs,
long nowMs) {
+ return isInFlight(tableId, policyId)
+ || 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;
+ }
+ JobEntity job = loadJob(open.jobId(), tableId);
+ if (job == null || ACTIVE.contains(job.status())) {
+ continue;
+ }
+ Long fromJob = job.finishedAt() != null && job.finishedAt() > 0 ?
job.finishedAt() : null;
+ if (fromJob == null) {
+ continue;
+ }
+ if (latest == null || fromJob > latest) {
+ latest = fromJob;
+ }
+ }
+ return latest;
+ }
+
+ @Nullable
+ private static JobEntity loadJob(long jobId, long tableId) {
+ String metalake = metalakeNameForTable(tableId);
+ if (metalake == null) {
+ return null;
+ }
+ try {
+ return JobMetaService.getInstance()
+ .getJobByIdentifier(NameIdentifierUtil.ofJob(metalake,
JobHandle.JOB_ID_PREFIX + jobId));
+ } catch (NoSuchEntityException e) {
+ return null;
+ }
+ }
+
+ @Nullable
+ private static String metalakeNameForTable(long tableId) {
+ List<TablePO> tables =
+ SessionUtils.getWithoutCommit(
+ TableMetaMapper.class,
+ mapper ->
mapper.listTablePOsByTableIds(Collections.singletonList(tableId)));
+ if (tables == null || tables.isEmpty()) {
+ return null;
+ }
+ MetalakePO metalake =
+ SessionUtils.getWithoutCommit(
+ MetalakeMetaMapper.class,
+ mapper ->
mapper.selectMetalakeMetaById(tables.get(0).getMetalakeId()));
+ return metalake == null ? null : metalake.getMetalakeName();
+ }
+}
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 419a1987fb..4ced139557 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
@@ -146,7 +146,8 @@ public class TableMaintenanceActionBaseSQLProvider {
}
/**
- * Sets {@code before_snapshot_id} once when {@code S0} is sampled.
+ * 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 beforeSnapshotId sampled snapshot id
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 9814cc3c34..ae1cd03e11 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
@@ -216,7 +216,10 @@ public class TableMaintenanceActionMetaService {
}
/**
- * Records the sampled Iceberg snapshot id ({@code S0}) once for a running
action.
+ * Records the sampled Iceberg snapshot id ({@code S0}) for a running action.
+ *
+ * <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 beforeSnapshotId sampled snapshot id
@@ -229,10 +232,19 @@ public class TableMaintenanceActionMetaService {
SessionUtils.doWithCommitAndFetchResult(
TableMaintenanceActionMapper.class,
mapper -> mapper.updateBeforeSnapshot(jobId, beforeSnapshotId));
+ if (updated == 1) {
+ return;
+ }
+ TableMaintenanceActionPO current = getByJobId(jobId);
Preconditions.checkState(
- updated == 1,
- "expected to set before_snapshot_id once on an unfinished action,
updated=%s",
- updated);
+ current != null && current.finishedAt() != null &&
current.finishedAt() == 0L,
+ "expected an unfinished action for job_id=%s",
+ jobId);
+ Preconditions.checkState(
+ current.beforeSnapshotId() != null && current.beforeSnapshotId() ==
beforeSnapshotId,
+ "before_snapshot_id already set to %s, refusing %s",
+ current.beforeSnapshotId(),
+ beforeSnapshotId);
}
/**
@@ -302,6 +314,29 @@ public class TableMaintenanceActionMetaService {
return open == null ? Collections.emptyList() : open;
}
+ /**
+ * Whether an unbound claim is still within TTL on the
<strong>database</strong> clock.
+ *
+ * <p>Must not be evaluated with {@code System.currentTimeMillis()} against
{@code created_at}:
+ * acquire/reclaim already use DB time, and JVM skew would disagree with
that lease.
+ *
+ * @param tableId table entity id
+ * @param policyId policy entity id
+ * @param claimTimeoutMs unbound claim TTL
+ * @return true when a live unbound claim holds the soft slot
+ */
+ @Monitored(
+ metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME,
+ baseMetricName = "hasLiveUnboundTableMaintenanceAction")
+ public boolean hasLiveUnboundClaim(long tableId, long policyId, long
claimTimeoutMs) {
+ Preconditions.checkArgument(claimTimeoutMs > 0, "claimTimeoutMs must be
positive");
+ TableMaintenanceActionPO live =
+ SessionUtils.getWithoutCommit(
+ TableMaintenanceActionMapper.class,
+ mapper -> mapper.selectLiveUnboundByTableAndPolicy(tableId,
policyId, claimTimeoutMs));
+ return live != null;
+ }
+
/**
* @param tableId table entity id
* @param policyId policy entity id
diff --git
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableSnapshotMetricsMetaService.java
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableSnapshotMetricsMetaService.java
index 392c846b94..fe7b0719d2 100644
---
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableSnapshotMetricsMetaService.java
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableSnapshotMetricsMetaService.java
@@ -20,6 +20,8 @@ package org.apache.gravitino.maintenance.storage.service;
import static
org.apache.gravitino.metrics.source.MetricsSource.GRAVITINO_RELATIONAL_STORE_METRIC_NAME;
+import java.sql.SQLIntegrityConstraintViolationException;
+import java.util.Locale;
import javax.annotation.Nullable;
import
org.apache.gravitino.maintenance.storage.mapper.TableSnapshotMetricsMapper;
import org.apache.gravitino.maintenance.storage.po.TableSnapshotMetricsPO;
@@ -45,6 +47,9 @@ public class TableSnapshotMetricsMetaService {
* Idempotently stores metrics JSON for a table/policy/snapshot triple. If a
row already exists,
* this is a no-op (does not overwrite) so sample / listener retries are
safe.
*
+ * <p>A concurrent insert that loses the unique-key race is treated as
success (first writer
+ * wins).
+ *
* @param tableId table id
* @param policyId policy id
* @param snapshotId Iceberg snapshot id
@@ -65,8 +70,15 @@ public class TableSnapshotMetricsMetaService {
.withMetricsValue(metricsJson)
.withCreatedAt(System.currentTimeMillis())
.build();
- SessionUtils.doWithCommit(
- TableSnapshotMetricsMapper.class, mapper ->
mapper.insertMetrics(metricsMeta));
+ try {
+ SessionUtils.doWithCommit(
+ TableSnapshotMetricsMapper.class, mapper ->
mapper.insertMetrics(metricsMeta));
+ } catch (RuntimeException e) {
+ if (isUniqueViolation(e)) {
+ return;
+ }
+ throw e;
+ }
}
/**
@@ -84,4 +96,20 @@ public class TableSnapshotMetricsMetaService {
TableSnapshotMetricsMapper.class,
mapper -> mapper.selectMetrics(tableId, policyId, snapshotId));
}
+
+ private static boolean isUniqueViolation(Throwable throwable) {
+ for (Throwable t = throwable; t != null; t = t.getCause()) {
+ if (t instanceof SQLIntegrityConstraintViolationException) {
+ return true;
+ }
+ String message = t.getMessage();
+ if (message != null) {
+ String lower = message.toLowerCase(Locale.ROOT);
+ if (lower.contains("unique") || lower.contains("duplicate")) {
+ return true;
+ }
+ }
+ }
+ return false;
+ }
}
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
new file mode 100644
index 0000000000..ec9b9511bc
--- /dev/null
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/feature/TableMaintenanceRESTFeature.java
@@ -0,0 +1,130 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.feature;
+
+import java.util.concurrent.atomic.AtomicReference;
+import javax.sql.DataSource;
+import javax.ws.rs.core.Feature;
+import javax.ws.rs.core.FeatureContext;
+import javax.ws.rs.ext.Provider;
+import org.apache.gravitino.Config;
+import org.apache.gravitino.Configs;
+import org.apache.gravitino.GravitinoEnv;
+import org.apache.gravitino.maintenance.TableMaintenanceConfig;
+import
org.apache.gravitino.maintenance.scheduler.TableMaintenanceSchedulerPools;
+import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper;
+import org.glassfish.jersey.server.monitoring.ApplicationEvent;
+import org.glassfish.jersey.server.monitoring.ApplicationEventListener;
+import org.glassfish.jersey.server.monitoring.RequestEvent;
+import org.glassfish.jersey.server.monitoring.RequestEventListener;
+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).
+ *
+ * <p>TMS does not register public REST resources in this change; the feature
only owns server
+ * lifecycle for in-process scheduling. 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 and
+ * re-register the destroy listener (needed for tests / embedded restarts in
the same JVM).
+ */
+@Provider
+public class TableMaintenanceRESTFeature implements Feature {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(TableMaintenanceRESTFeature.class);
+
+ /** Package name for {@link Configs#REST_API_EXTENSION_PACKAGES}. */
+ public static final String EXTENSION_PACKAGE =
TableMaintenanceRESTFeature.class.getPackageName();
+
+ private static final AtomicReference<TableMaintenanceSchedulerPools>
SCHEDULER_POOLS =
+ new AtomicReference<>();
+
+ @Override
+ public boolean configure(FeatureContext context) {
+ if (SCHEDULER_POOLS.get() != null) {
+ LOG.warn("Table maintenance scheduler pools are already running;
skipping reconfigure");
+ return true;
+ }
+
+ Config config = GravitinoEnv.getInstance().config();
+ if
(!Configs.RELATIONAL_ENTITY_STORE.equals(config.get(Configs.ENTITY_STORE))) {
+ throw new IllegalStateException(
+ "Table maintenance requires gravitino.entity.store = relational");
+ }
+
+ DataSource dataSource = resolveEntityStoreDataSource();
+ TableMaintenanceSchedulerPools pools = new
TableMaintenanceSchedulerPools();
+ pools.start(dataSource, TableMaintenanceConfig.from(config));
+ if (!SCHEDULER_POOLS.compareAndSet(null, pools)) {
+ pools.stop();
+ LOG.warn("Table maintenance scheduler pools were started concurrently;
discarded duplicate");
+ return true;
+ }
+
+ // 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());
+ LOG.info("Table maintenance scheduler pools are running");
+ return true;
+ }
+
+ /** Stops scheduler pools. Invoked from Jersey destroy and tests. */
+ public static void stopSchedulers() {
+ TableMaintenanceSchedulerPools pools = SCHEDULER_POOLS.getAndSet(null);
+ if (pools != null) {
+ pools.stop();
+ LOG.info("Table maintenance scheduler pools stopped");
+ }
+ }
+
+ private static DataSource resolveEntityStoreDataSource() {
+ Object dataSource =
+ SqlSessionFactoryHelper.getInstance()
+ .getSqlSessionFactory()
+ .getConfiguration()
+ .getEnvironment()
+ .getDataSource();
+ if (!(dataSource instanceof DataSource)) {
+ throw new IllegalStateException("Entity store data source is not
available");
+ }
+ return (DataSource) dataSource;
+ }
+
+ /**
+ * Stops TMS scheduler pools when the Jersey application is destroyed (Jetty
stop), before the
+ * relational entity store shuts down.
+ */
+ static final class SchedulerShutdownListener implements
ApplicationEventListener {
+
+ @Override
+ public void onEvent(ApplicationEvent event) {
+ if (event.getType() == ApplicationEvent.Type.DESTROY_FINISHED) {
+ stopSchedulers();
+ }
+ }
+
+ @Override
+ public RequestEventListener onRequest(RequestEvent requestEvent) {
+ return null;
+ }
+ }
+}
diff --git
a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/injob/TestTableMaintenanceInJobPipeline.java
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/injob/TestTableMaintenanceInJobPipeline.java
new file mode 100644
index 0000000000..c4cfbe9e56
--- /dev/null
+++
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/injob/TestTableMaintenanceInJobPipeline.java
@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.injob;
+
+import org.apache.gravitino.maintenance.TableMaintenanceConfig;
+import
org.apache.gravitino.maintenance.storage.service.TableMaintenanceActionMetaService;
+import
org.apache.gravitino.maintenance.storage.service.TableSnapshotMetricsMetaService;
+import org.apache.gravitino.storage.relational.TestJDBCBackend;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestTemplate;
+
+class TestTableMaintenanceInJobPipeline extends TestJDBCBackend {
+
+ private static final String METRICS_JSON = "{\"mse\":1}";
+
+ private TableMaintenanceActionMetaService actionMetaService;
+ private TableMaintenanceInJobPipeline maintainPipeline;
+ private TableMaintenanceInJobPipeline skipPipeline;
+
+ @BeforeEach
+ void setUp() {
+ actionMetaService = TableMaintenanceActionMetaService.getInstance();
+ TableMaintenanceInJobSampler sampler =
+ new TableMaintenanceInJobSampler(
+ TableSnapshotMetricsMetaService.getInstance(), actionMetaService);
+ maintainPipeline =
+ new TableMaintenanceInJobPipeline(
+ sampler, (tableId, policyId, snapshotId, metricsJson) ->
RecommendDecision.MAINTAIN);
+ skipPipeline =
+ new TableMaintenanceInJobPipeline(
+ sampler, (tableId, policyId, snapshotId, metricsJson) ->
RecommendDecision.SKIP);
+ }
+
+ @TestTemplate
+ void testContinuesWhenNoDriftAndRecommenderMaintains() {
+ long jobId = 2001L;
+ insertBoundAction(jobId, 10L, 20L);
+ InJobGateOutcome outcome =
+ maintainPipeline.gateCronJob(jobId, 10L, 20L, 100L, METRICS_JSON, ()
-> 100L);
+ Assertions.assertTrue(outcome.continueMaintenance());
+ Assertions.assertNull(outcome.stopReason());
+ }
+
+ @TestTemplate
+ void testStopsOnDriftBeforeRecommend() {
+ long jobId = 2002L;
+ insertBoundAction(jobId, 11L, 21L);
+ InJobGateOutcome outcome =
+ maintainPipeline.gateCronJob(jobId, 11L, 21L, 200L, METRICS_JSON, ()
-> 201L);
+ Assertions.assertFalse(outcome.continueMaintenance());
+ Assertions.assertEquals(InJobGateOutcome.StopReason.SNAPSHOT_DRIFT,
outcome.stopReason());
+ }
+
+ @TestTemplate
+ void testStopsWhenRecommenderSkips() {
+ long jobId = 2003L;
+ insertBoundAction(jobId, 12L, 22L);
+ InJobGateOutcome outcome =
+ skipPipeline.gateCronJob(jobId, 12L, 22L, 300L, METRICS_JSON, () ->
300L);
+ Assertions.assertFalse(outcome.continueMaintenance());
+ Assertions.assertEquals(InJobGateOutcome.StopReason.RECOMMENDER_SKIP,
outcome.stopReason());
+ }
+
+ @TestTemplate
+ void testCommitChainTypeRecommendGate() {
+ InJobGateOutcome maintain =
+ maintainPipeline.gateCommitChainType(1L, 2L, 400L, METRICS_JSON, () ->
400L);
+ Assertions.assertTrue(maintain.continueMaintenance());
+
+ InJobGateOutcome skip =
+ skipPipeline.gateCommitChainType(1L, 2L, 401L, METRICS_JSON, () ->
401L);
+ Assertions.assertEquals(InJobGateOutcome.StopReason.RECOMMENDER_SKIP,
skip.stopReason());
+ }
+
+ private void insertBoundAction(long jobId, long tableId, long policyId) {
+ actionMetaService.bindJobId(
+ actionMetaService.insertInFlight(
+ tableId, policyId,
TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS),
+ jobId);
+ }
+}
diff --git
a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/injob/TestTableMaintenanceInJobSampler.java
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/injob/TestTableMaintenanceInJobSampler.java
new file mode 100644
index 0000000000..8be445761a
--- /dev/null
+++
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/injob/TestTableMaintenanceInJobSampler.java
@@ -0,0 +1,107 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.injob;
+
+import org.apache.gravitino.maintenance.TableMaintenanceConfig;
+import org.apache.gravitino.maintenance.storage.po.TableSnapshotMetricsPO;
+import
org.apache.gravitino.maintenance.storage.service.TableMaintenanceActionMetaService;
+import
org.apache.gravitino.maintenance.storage.service.TableSnapshotMetricsMetaService;
+import org.apache.gravitino.storage.relational.TestJDBCBackend;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.TestTemplate;
+
+class TestTableMaintenanceInJobSampler extends TestJDBCBackend {
+
+ private static final String METRICS_JSON = "{\"mse\":3}";
+
+ private TableMaintenanceInJobSampler sampler;
+ private TableMaintenanceActionMetaService actionMetaService;
+ private TableSnapshotMetricsMetaService metricsMetaService;
+
+ @BeforeEach
+ void setUp() {
+ actionMetaService = TableMaintenanceActionMetaService.getInstance();
+ metricsMetaService = TableSnapshotMetricsMetaService.getInstance();
+ sampler = new TableMaintenanceInJobSampler(metricsMetaService,
actionMetaService);
+ }
+
+ @TestTemplate
+ void testCronJobSamplePersistsAndContinuesWhenNoDrift() {
+ long jobId = 1001L;
+ insertBoundAction(jobId, 1L, 2L);
+ InJobGateOutcome outcome =
+ sampler.sampleAndGateCronJob(jobId, 1L, 2L, 100L, METRICS_JSON, () ->
100L);
+
+ Assertions.assertTrue(outcome.continueMaintenance());
+ Assertions.assertNull(outcome.stopReason());
+ TableSnapshotMetricsPO stored = metricsMetaService.getMetrics(1L, 2L,
100L);
+ Assertions.assertNotNull(stored);
+ Assertions.assertEquals(METRICS_JSON, stored.metricsValue());
+ Assertions.assertEquals(100L,
actionMetaService.getByJobId(jobId).beforeSnapshotId());
+ }
+
+ @TestTemplate
+ void testCronJobEndsEarlyWhenSnapshotDriftsAfterSample() {
+ long jobId = 1002L;
+ insertBoundAction(jobId, 3L, 4L);
+ InJobGateOutcome outcome =
+ sampler.sampleAndGateCronJob(jobId, 3L, 4L, 200L, METRICS_JSON, () ->
201L);
+
+ Assertions.assertFalse(outcome.continueMaintenance());
+ Assertions.assertEquals(InJobGateOutcome.StopReason.SNAPSHOT_DRIFT,
outcome.stopReason());
+ Assertions.assertEquals(200L, outcome.sampledSnapshotId());
+ Assertions.assertEquals(201L, outcome.headSnapshotId());
+ Assertions.assertNotNull(metricsMetaService.getMetrics(3L, 4L, 200L));
+ }
+
+ @TestTemplate
+ void testCronJobSampleRetryWithSameS0IsIdempotent() {
+ long jobId = 1003L;
+ insertBoundAction(jobId, 5L, 6L);
+ InJobGateOutcome first =
+ sampler.sampleAndGateCronJob(jobId, 5L, 6L, 300L, METRICS_JSON, () ->
300L);
+ InJobGateOutcome retry =
+ sampler.sampleAndGateCronJob(jobId, 5L, 6L, 300L, "{\"mse\":9}", () ->
300L);
+
+ Assertions.assertTrue(first.continueMaintenance());
+ Assertions.assertTrue(retry.continueMaintenance());
+ Assertions.assertEquals(
+ METRICS_JSON, metricsMetaService.getMetrics(5L, 6L,
300L).metricsValue());
+ Assertions.assertEquals(300L,
actionMetaService.getByJobId(jobId).beforeSnapshotId());
+ }
+
+ @TestTemplate
+ void testCommitChainTypeGateWithoutReSample() {
+ InJobGateOutcome ok = sampler.gateCommitChainType(300L, () -> 300L);
+ Assertions.assertTrue(ok.continueMaintenance());
+
+ InJobGateOutcome drift = sampler.gateCommitChainType(300L, () -> 301L);
+ Assertions.assertFalse(drift.continueMaintenance());
+ Assertions.assertEquals(InJobGateOutcome.StopReason.SNAPSHOT_DRIFT,
drift.stopReason());
+ }
+
+ private void insertBoundAction(long jobId, long tableId, long policyId) {
+ actionMetaService.bindJobId(
+ actionMetaService.insertInFlight(
+ tableId, policyId,
TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS),
+ jobId);
+ }
+}
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
new file mode 100644
index 0000000000..f63532af35
--- /dev/null
+++
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/occupancy/TestOccupancyGate.java
@@ -0,0 +1,221 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.occupancy;
+
+import java.io.IOException;
+import java.util.function.Function;
+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.TableEntity;
+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 testInFlightReadsJobMetaServiceWithoutBackfillingFinishedAt() 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);
+ long policyId = 22L;
+
+ Assertions.assertFalse(gate.isInFlight(table.id(), policyId));
+
+ OccupancyClaim claim =
+ actions.insertInFlight(
+ 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());
+
+ 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());
+
+ Assertions.assertFalse(gate.isInFlight(table.id(), policyId));
+ Assertions.assertEquals(0L, 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));
+ }
+
+ @TestTemplate
+ void testMinIntervalUsesMarkedFinishedAt() throws IOException {
+ String metalakeName = "tms_occ_ml2";
+ 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);
+ long policyId = 44L;
+ long jobId = 7002L;
+ long finishedAt = 1_700_000_100_000L;
+
+ OccupancyClaim claim =
+ actions.insertInFlight(
+ table.id(), policyId,
TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS);
+ actions.bindJobId(claim, jobId);
+ actions.markFinished(jobId, 200L, finishedAt);
+ Assertions.assertFalse(gate.isInFlight(table.id(), policyId));
+ Assertions.assertTrue(
+ gate.withinMinInterval(table.id(), policyId, 3_600_000L, finishedAt +
1_000L));
+ }
+
+ @TestTemplate
+ void testSoftSlotBlocksNewAcquireUntilMarkFinished() throws IOException {
+ 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 policyId = 55L;
+ 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.assertThrows(
+ IllegalStateException.class,
+ () ->
+ actions.insertInFlight(
+ table.id(), policyId,
TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS));
+
+ actions.markFinished(succeeded.id(), 200L, 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());
+ Assertions.assertTrue(gate.isInFlight(table.id(), policyId));
+ Assertions.assertTrue(
+ gate.shouldSkip(table.id(), policyId, 3_600_000L, finishedAt +
3_600_001L));
+ }
+
+ @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 = 200L;
+ 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));
+
+ Thread.sleep(claimTimeoutMs + 50L);
+
+ // 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());
+ }
+}
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 855b073d49..5979eef5d8 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
@@ -50,6 +50,9 @@ class TestTableMaintenanceActionMetaService extends
TestJDBCBackend {
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.assertThrows(
diff --git
a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/web/rest/feature/TestTableMaintenanceRESTFeature.java
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/web/rest/feature/TestTableMaintenanceRESTFeature.java
new file mode 100644
index 0000000000..bd90aa2f77
--- /dev/null
+++
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/web/rest/feature/TestTableMaintenanceRESTFeature.java
@@ -0,0 +1,40 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.gravitino.maintenance.web.rest.feature;
+
+import org.glassfish.jersey.server.monitoring.ApplicationEvent;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.Test;
+import org.mockito.Mockito;
+
+class TestTableMaintenanceRESTFeature {
+
+ @Test
+ void testDestroyFinishedStopsSchedulers() {
+ TableMaintenanceRESTFeature.stopSchedulers();
+ ApplicationEvent event = Mockito.mock(ApplicationEvent.class);
+
Mockito.when(event.getType()).thenReturn(ApplicationEvent.Type.DESTROY_FINISHED);
+
+ TableMaintenanceRESTFeature.SchedulerShutdownListener listener =
+ new TableMaintenanceRESTFeature.SchedulerShutdownListener();
+ Assertions.assertDoesNotThrow(() -> listener.onEvent(event));
+ // stop is idempotent so configure-after-stop can start pools again and
re-register a listener
+ Assertions.assertDoesNotThrow(TableMaintenanceRESTFeature::stopSchedulers);
+ }
+}
