This is an automated email from the ASF dual-hosted git repository.
roryqi pushed a commit to branch branch-tms-scheduler
in repository https://gitbox.apache.org/repos/asf/gravitino.git
The following commit(s) were added to refs/heads/branch-tms-scheduler by this
push:
new 45a188f071 [#13983] feat(tms): Register policy timeTrigger on the
expand scheduler (#13897)
45a188f071 is described below
commit 45a188f071531f136f6c15fc86d5540944963a2a
Author: MaSai <[email protected]>
AuthorDate: Fri Oct 23 17:58:58 2026 +0800
[#13983] feat(tms): Register policy timeTrigger on the expand scheduler
(#13897)
### What changes were proposed in this pull request?
Follow-up on the TMS scheduler pools from Epic #13760.
When a maintenance policy carries `timeTrigger`, create, alter, enable,
disable, and delete keep a matching `policy-expand` row in
`scheduled_expand_tasks`:
- Enabled policy with a non-blank cron: insert or move the row. The task
instance id is the policy id, and the due time is the next Quartz cron
fire time.
- Disabled policy, blank `timeTrigger`, or a successful delete: remove
the row.
- Non-maintenance policies are ignored.
The expand handler is still a no-op. This change only registers the row.
The db-scheduler client forces auto-commit, and the pools commit when
auto-commit is off, so writes survive the entity-store connection pool.
### Why are the changes needed?
Policy content stores `timeTrigger`, but nothing was writing that
schedule into db-scheduler. Without this, a created or updated
maintenance policy never becomes an expand lease, and a deleted or
disabled policy would leave a stale row.
Fix: #13983
Parent epic: #13760
### Does this PR introduce _any_ user-facing change?
No new REST API or configuration key. Creating, altering, enabling,
disabling, or deleting a maintenance policy that has `timeTrigger` now
inserts, moves, or deletes the corresponding `policy-expand` row while
TMS scheduler pools are running.
### How was this patch tested?
- `TestPolicyTimeTriggerEventListener` covers create, alter, enable,
disable, delete, a failed delete, and non-maintenance policies.
- `TestPolicyExpandScheduler` inserts, reschedules, and removes the
expand row on the H2 entity-store schema.
- `./gradlew :plugins:tms:test --tests
org.apache.gravitino.maintenance.scheduler.TestPolicyTimeTriggerEventListener
--tests
org.apache.gravitino.maintenance.scheduler.TestPolicyExpandScheduler
--tests
org.apache.gravitino.maintenance.scheduler.TestTableMaintenanceSchedulerPools
-PskipITs`
Made with [Cursor](https://cursor.com)
---------
Co-authored-by: Cursor <[email protected]>
---
.../gravitino/policy/hook/PolicyLifecycleHook.java | 88 ++++++
.../policy/hook/PolicyLifecycleHooks.java | 66 +++++
.../relational/service/PolicyMetaService.java | 67 ++++-
.../storage/relational/session/SqlSessions.java | 25 +-
.../policy/hook/TestPolicyLifecycleHooks.java | 59 ++++
.../relational/service/TestPolicyMetaService.java | 156 ++++++++++
.../policy/PolicyExpandLifecycleHook.java | 203 +++++++++++++
.../maintenance/policy/PolicyExpandReconciler.java | 118 ++++++++
.../maintenance/policy/PolicyExpandRegistrar.java | 325 +++++++++++++++++++++
.../scheduler/TableMaintenanceSchedulerPools.java | 37 ++-
.../rest/feature/TableMaintenanceRESTFeature.java | 58 +++-
.../policy/TestPolicyExpandLifecycleHook.java | 135 +++++++++
.../policy/TestPolicyExpandReconciler.java | 249 ++++++++++++++++
.../policy/TestPolicyExpandRegistrar.java | 135 +++++++++
.../TestTableMaintenanceSchedulerPools.java | 2 +-
15 files changed, 1704 insertions(+), 19 deletions(-)
diff --git
a/core/src/main/java/org/apache/gravitino/policy/hook/PolicyLifecycleHook.java
b/core/src/main/java/org/apache/gravitino/policy/hook/PolicyLifecycleHook.java
new file mode 100644
index 0000000000..f8f03bf5f6
--- /dev/null
+++
b/core/src/main/java/org/apache/gravitino/policy/hook/PolicyLifecycleHook.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.policy.hook;
+
+import org.apache.gravitino.meta.PolicyEntity;
+
+/**
+ * Storage-layer callbacks inside policy create / update / delete
transactions, with JPA-style
+ * lifecycle points relative to the policy DML.
+ *
+ * <p>Distinct from {@code org.apache.gravitino.hook.PolicyHookDispatcher},
which decorates the
+ * policy API after a successful manager call (for example ownership). This
SPI runs in {@code
+ * PolicyMetaService} on the open SQL session so extension writes commit or
roll back with the
+ * policy row.
+ *
+ * <p>Implementations must use that connection, and must not commit, roll
back, or close it. An
+ * exception rolls the policy write back together with the hook's own writes.
+ *
+ * <p>Defaults are no-ops so an extension only overrides the points it needs.
+ *
+ * <p>{@code insertPolicy(..., overwritten = true)} that replaces an existing
row fires the update
+ * lifecycle ({@link #beforeUpdate} / {@link #afterUpdate}), not persist. A
create, or an overwrite
+ * that inserts because no row was found, fires persist.
+ */
+public interface PolicyLifecycleHook {
+
+ /**
+ * About to insert the policy row in the current transaction.
+ *
+ * @param policy the policy about to be stored
+ */
+ default void beforePersist(PolicyEntity policy) {}
+
+ /**
+ * The policy row has been inserted in the current transaction.
+ *
+ * @param policy the policy as stored, including id, enabled flag, and
content
+ */
+ default void afterPersist(PolicyEntity policy) {}
+
+ /**
+ * About to update the policy row in the current transaction.
+ *
+ * <p>Also used when {@code insertPolicy} overwrites an existing row by name
or id.
+ *
+ * @param policy the policy about to be stored (id is the retained row id on
overwrite)
+ */
+ default void beforeUpdate(PolicyEntity policy) {}
+
+ /**
+ * The policy row has been updated in the current transaction.
+ *
+ * <p>Also used when {@code insertPolicy} overwrites an existing row by name
or id.
+ *
+ * @param policy the policy as stored, including the retained id, enabled
flag, and content
+ */
+ default void afterUpdate(PolicyEntity policy) {}
+
+ /**
+ * The policy row is about to be deleted in the current transaction.
+ *
+ * @param policyId persisted policy id
+ */
+ default void beforeRemove(long policyId) {}
+
+ /**
+ * The policy row has been deleted in the current transaction.
+ *
+ * @param policyId persisted policy id
+ */
+ default void afterRemove(long policyId) {}
+}
diff --git
a/core/src/main/java/org/apache/gravitino/policy/hook/PolicyLifecycleHooks.java
b/core/src/main/java/org/apache/gravitino/policy/hook/PolicyLifecycleHooks.java
new file mode 100644
index 0000000000..9e6f7a9c3c
--- /dev/null
+++
b/core/src/main/java/org/apache/gravitino/policy/hook/PolicyLifecycleHooks.java
@@ -0,0 +1,66 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.gravitino.policy.hook;
+
+import java.util.concurrent.atomic.AtomicReference;
+import javax.annotation.Nullable;
+
+/**
+ * Holds the optional {@link PolicyLifecycleHook} registered by an extension.
+ *
+ * <p>This is a storage-layer SPI invoked from {@code PolicyMetaService}
inside the open policy
+ * transaction. It is not the same as {@code
org.apache.gravitino.hook.PolicyHookDispatcher}, which
+ * wraps the policy API for owner side effects after a successful manager call.
+ */
+public final class PolicyLifecycleHooks {
+
+ private static final AtomicReference<PolicyLifecycleHook> HOOK = new
AtomicReference<>();
+
+ private PolicyLifecycleHooks() {}
+
+ /**
+ * Installs the hook invoked at policy create / update / delete lifecycle
points inside the open
+ * transaction.
+ *
+ * @param hook hook to run during policy create, update, and delete
transactions
+ * @throws IllegalStateException if a hook is already registered
+ */
+ public static void register(PolicyLifecycleHook hook) {
+ if (!HOOK.compareAndSet(null, hook)) {
+ throw new IllegalStateException("A policy lifecycle hook is already
registered");
+ }
+ }
+
+ /**
+ * Removes the hook when it is the one currently registered.
+ *
+ * @param hook hook passed to {@link #register}
+ */
+ public static void unregister(PolicyLifecycleHook hook) {
+ HOOK.compareAndSet(hook, null);
+ }
+
+ /**
+ * @return the registered hook, or {@code null} when no extension is
participating
+ */
+ @Nullable
+ public static PolicyLifecycleHook get() {
+ return HOOK.get();
+ }
+}
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/service/PolicyMetaService.java
b/core/src/main/java/org/apache/gravitino/storage/relational/service/PolicyMetaService.java
index 4d3a1abc39..0327b97bfb 100644
---
a/core/src/main/java/org/apache/gravitino/storage/relational/service/PolicyMetaService.java
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/service/PolicyMetaService.java
@@ -40,6 +40,8 @@ import org.apache.gravitino.Namespace;
import org.apache.gravitino.exceptions.NoSuchEntityException;
import org.apache.gravitino.meta.PolicyEntity;
import org.apache.gravitino.metrics.Monitored;
+import org.apache.gravitino.policy.hook.PolicyLifecycleHook;
+import org.apache.gravitino.policy.hook.PolicyLifecycleHooks;
import org.apache.gravitino.storage.relational.mapper.MetalakeMetaMapper;
import org.apache.gravitino.storage.relational.mapper.OwnerMetaMapper;
import org.apache.gravitino.storage.relational.mapper.PolicyMetaMapper;
@@ -172,8 +174,10 @@ public class PolicyMetaService {
PolicyPO newPolicyPO =
POConverters.updatePolicyPOWithVersion(oldPolicyPO,
updatedPolicyEntity);
SessionUtils.doMultipleWithCommit(
+ () -> beforeUpdate(updatedPolicyEntity),
() -> updatePolicyRootWithVersion(ident, oldPolicyPO, newPolicyPO),
- () -> insertPolicyVersionIfAllocated(oldPolicyPO, newPolicyPO));
+ () -> insertPolicyVersionIfAllocated(oldPolicyPO, newPolicyPO),
+ () -> afterUpdate(updatedPolicyEntity));
} catch (RuntimeException re) {
ExceptionUtils.checkSQLException(
re, Entity.EntityType.POLICY,
updatedPolicyEntity.nameIdentifier().toString());
@@ -212,6 +216,7 @@ public class PolicyMetaService {
try {
SessionUtils.doMultipleWithCommit(
+ () -> beforeRemove(policyId),
() -> deletePolicyWithVersion(ident, policyPO),
() ->
SessionUtils.doWithoutCommit(
@@ -237,7 +242,8 @@ public class PolicyMetaService {
SecurableObjectMapper.class,
mapper ->
mapper.softDeleteObjectRelsByMetadataObject(
- policyId, MetadataObject.Type.POLICY.name())));
+ policyId, MetadataObject.Type.POLICY.name())),
+ () -> afterRemove(policyId));
return true;
} catch (NoSuchEntityException e) {
return false;
@@ -363,7 +369,12 @@ public class PolicyMetaService {
}
/**
- * Writes the policy row and its content snapshot.
+ * Writes the policy row and its content snapshot without committing, and
fires the matching
+ * lifecycle hooks around the DML.
+ *
+ * <p>A plain create (or an overwrite that finds no row) uses persist hooks.
An overwrite that
+ * replaces an existing row uses update hooks and passes the entity as
stored — including the
+ * retained policy id when the caller supplied a different id for the same
name.
*
* <p>An overwrite is not an upsert any more: the existing row is located
and locked first, and
* the replacement is written as the next version of that row, so the
snapshot history survives
@@ -378,7 +389,9 @@ public class PolicyMetaService {
private void insertPolicyWithoutCommit(
PolicyEntity policyEntity, PolicyPO initializedPolicyPO, boolean
overwritten) {
if (!overwritten) {
+ beforePersist(policyEntity);
insertNewPolicyWithoutCommit(initializedPolicyPO);
+ afterPersist(policyEntity);
return;
}
@@ -392,16 +405,22 @@ public class PolicyMetaService {
"The policy ID %s is reserved by a deleted policy; use a new ID",
initializedPolicyPO.getPolicyId());
}
+ beforePersist(policyEntity);
insertNewPolicyWithoutCommit(initializedPolicyPO);
+ afterPersist(policyEntity);
return;
}
PolicyPO replacementPolicyPO =
POConverters.updatePolicyPOWithVersion(existingPolicyPO,
initializedPolicyPO);
+ PolicyEntity storedPolicy =
+ POConverters.fromPolicyPO(replacementPolicyPO,
policyEntity.namespace());
NameIdentifier observedIdentifier =
NameIdentifier.of(policyEntity.namespace(),
existingPolicyPO.getPolicyName());
+ beforeUpdate(storedPolicy);
updatePolicyRootWithVersion(observedIdentifier, existingPolicyPO,
replacementPolicyPO);
insertPolicyVersionIfAllocated(existingPolicyPO, replacementPolicyPO);
+ afterUpdate(storedPolicy);
}
/**
@@ -570,4 +589,46 @@ public class PolicyMetaService {
}
return policyPO;
}
+
+ private static void beforePersist(PolicyEntity policy) {
+ PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+ if (hook != null) {
+ hook.beforePersist(policy);
+ }
+ }
+
+ private static void afterPersist(PolicyEntity policy) {
+ PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+ if (hook != null) {
+ hook.afterPersist(policy);
+ }
+ }
+
+ private static void beforeUpdate(PolicyEntity policy) {
+ PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+ if (hook != null) {
+ hook.beforeUpdate(policy);
+ }
+ }
+
+ private static void afterUpdate(PolicyEntity policy) {
+ PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+ if (hook != null) {
+ hook.afterUpdate(policy);
+ }
+ }
+
+ private static void beforeRemove(long policyId) {
+ PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+ if (hook != null) {
+ hook.beforeRemove(policyId);
+ }
+ }
+
+ private static void afterRemove(long policyId) {
+ PolicyLifecycleHook hook = PolicyLifecycleHooks.get();
+ if (hook != null) {
+ hook.afterRemove(policyId);
+ }
+ }
}
diff --git
a/core/src/main/java/org/apache/gravitino/storage/relational/session/SqlSessions.java
b/core/src/main/java/org/apache/gravitino/storage/relational/session/SqlSessions.java
index 48a10324b5..2ebe3d30b4 100644
---
a/core/src/main/java/org/apache/gravitino/storage/relational/session/SqlSessions.java
+++
b/core/src/main/java/org/apache/gravitino/storage/relational/session/SqlSessions.java
@@ -20,6 +20,8 @@
package org.apache.gravitino.storage.relational.session;
import com.google.common.annotations.VisibleForTesting;
+import java.sql.Connection;
+import java.sql.SQLException;
import java.util.concurrent.atomic.AtomicInteger;
import org.apache.ibatis.session.SqlSession;
import org.apache.ibatis.session.TransactionIsolationLevel;
@@ -112,6 +114,23 @@ public final class SqlSessions {
return sessions.get() != null;
}
+ /**
+ * Returns the connection of the current thread's open session without
changing the session count.
+ *
+ * <p>Callers must not commit, roll back, or close the connection. The
session owner commits it.
+ *
+ * @return the connection bound to the current session
+ * @throws IllegalStateException if the current thread has no open session
+ * @throws SQLException if the session cannot provide its connection
+ */
+ public static Connection currentConnection() throws SQLException {
+ SqlSession sqlSession = sessions.get();
+ if (sqlSession == null) {
+ throw new IllegalStateException("No active SQL session on the current
thread");
+ }
+ return sqlSession.getConnection();
+ }
+
private static void handleSessionClose(boolean commit, boolean rollback) {
SqlSession sqlSession = sessions.get();
if (sqlSession == null) {
@@ -124,7 +143,11 @@ public final class SqlSessions {
if (commit) {
sqlSession.commit();
} else if (rollback) {
- sqlSession.rollback();
+ // Force the JDBC rollback. A policy lifecycle hook can write on
this connection without
+ // going
+ // through a mapped statement, and MyBatis skips rollback when the
session looks clean.
+ // Closing that session then turns auto-commit back on, which would
commit the hook write.
+ sqlSession.rollback(true);
}
} finally {
try {
diff --git
a/core/src/test/java/org/apache/gravitino/policy/hook/TestPolicyLifecycleHooks.java
b/core/src/test/java/org/apache/gravitino/policy/hook/TestPolicyLifecycleHooks.java
new file mode 100644
index 0000000000..6d802b9167
--- /dev/null
+++
b/core/src/test/java/org/apache/gravitino/policy/hook/TestPolicyLifecycleHooks.java
@@ -0,0 +1,59 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.gravitino.policy.hook;
+
+import static org.junit.jupiter.api.Assertions.assertNull;
+import static org.junit.jupiter.api.Assertions.assertSame;
+import static org.junit.jupiter.api.Assertions.assertThrows;
+
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.Test;
+
+public class TestPolicyLifecycleHooks {
+
+ private final PolicyLifecycleHook hook = new PolicyLifecycleHook() {};
+
+ @AfterEach
+ void tearDown() {
+ PolicyLifecycleHooks.unregister(hook);
+ }
+
+ @Test
+ void testRegisterAndUnregister() {
+ PolicyLifecycleHooks.register(hook);
+ assertSame(hook, PolicyLifecycleHooks.get());
+ PolicyLifecycleHooks.unregister(hook);
+ assertNull(PolicyLifecycleHooks.get());
+ }
+
+ @Test
+ void testDoubleRegisterRejected() {
+ PolicyLifecycleHooks.register(hook);
+ PolicyLifecycleHook other = new PolicyLifecycleHook() {};
+ assertThrows(IllegalStateException.class, () ->
PolicyLifecycleHooks.register(other));
+ assertSame(hook, PolicyLifecycleHooks.get());
+ }
+
+ @Test
+ void testUnregisterWrongHookIsNoOp() {
+ PolicyLifecycleHooks.register(hook);
+ PolicyLifecycleHooks.unregister(new PolicyLifecycleHook() {});
+ assertSame(hook, PolicyLifecycleHooks.get());
+ }
+}
diff --git
a/core/src/test/java/org/apache/gravitino/storage/relational/service/TestPolicyMetaService.java
b/core/src/test/java/org/apache/gravitino/storage/relational/service/TestPolicyMetaService.java
index 74dc33e67b..8a89cc9014 100644
---
a/core/src/test/java/org/apache/gravitino/storage/relational/service/TestPolicyMetaService.java
+++
b/core/src/test/java/org/apache/gravitino/storage/relational/service/TestPolicyMetaService.java
@@ -34,6 +34,7 @@ import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.time.Instant;
+import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
@@ -65,6 +66,8 @@ import org.apache.gravitino.meta.UserEntity;
import org.apache.gravitino.policy.Policy;
import org.apache.gravitino.policy.PolicyContent;
import org.apache.gravitino.policy.PolicyContents;
+import org.apache.gravitino.policy.hook.PolicyLifecycleHook;
+import org.apache.gravitino.policy.hook.PolicyLifecycleHooks;
import org.apache.gravitino.storage.RandomIdGenerator;
import org.apache.gravitino.storage.relational.TestJDBCBackend;
import org.apache.gravitino.storage.relational.mapper.PolicyMetaMapper;
@@ -93,6 +96,159 @@ public class TestPolicyMetaService extends TestJDBCBackend {
private final PolicyContent content =
PolicyContents.custom(ImmutableMap.of("filed1", 123),
supportedObjectTypes, null);
+ @TestTemplate
+ public void testLifecycleHookFailureRollsBackInsert() throws IOException {
+ BaseMetalake metalake = createAndInsertMakeLake(METALAKE_NAME);
+ PolicyLifecycleHook hook =
+ new PolicyLifecycleHook() {
+ @Override
+ public void afterPersist(PolicyEntity policy) {
+ throw new IllegalArgumentException("bad cron");
+ }
+ };
+ PolicyLifecycleHooks.register(hook);
+ try {
+ PolicyEntity policyEntity =
+ PolicyEntity.builder()
+ .withId(RandomIdGenerator.INSTANCE.nextId())
+ .withName("hook-rollback")
+ .withNamespace(NamespaceUtil.ofPolicy(metalake.name()))
+ .withComment("comment")
+ .withPolicyType(Policy.BuiltInType.CUSTOM)
+ .withContent(content)
+ .withEnabled(true)
+ .withAuditInfo(AUDIT_INFO)
+ .build();
+ PolicyMetaService policyMetaService = PolicyMetaService.getInstance();
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> policyMetaService.insertPolicy(policyEntity, false));
+ assertThrows(
+ NoSuchEntityException.class,
+ () ->
+ policyMetaService.getPolicyByIdentifier(
+ NameIdentifierUtil.ofPolicy(metalake.name(),
"hook-rollback")));
+ } finally {
+ PolicyLifecycleHooks.unregister(hook);
+ }
+ }
+
+ @TestTemplate
+ public void testLifecycleHookFailureRollsBackUpdate() throws IOException {
+ createAndInsertMakeLake(METALAKE_NAME);
+ PolicyMetaService policyMetaService = PolicyMetaService.getInstance();
+ PolicyEntity policy =
+ createPolicy(
+ RandomIdGenerator.INSTANCE.nextId(),
+ NamespaceUtil.ofPolicy(METALAKE_NAME),
+ "hook-update-rollback",
+ AUDIT_INFO);
+ policyMetaService.insertPolicy(policy, false);
+
+ PolicyLifecycleHook hook =
+ new PolicyLifecycleHook() {
+ @Override
+ public void afterUpdate(PolicyEntity updated) {
+ throw new IllegalArgumentException("bad cron on update");
+ }
+ };
+ PolicyLifecycleHooks.register(hook);
+ try {
+ assertThrows(
+ IllegalArgumentException.class,
+ () ->
+ policyMetaService.updatePolicy(
+ policy.nameIdentifier(),
+ entity -> copyPolicy((PolicyEntity) entity, policy.name(),
"should-not-stick")));
+ assertEquals(
+ policy.comment(),
+
policyMetaService.getPolicyByIdentifier(policy.nameIdentifier()).comment());
+ } finally {
+ PolicyLifecycleHooks.unregister(hook);
+ }
+ }
+
+ @TestTemplate
+ public void testLifecycleHookFailureRollsBackDelete() throws IOException {
+ createAndInsertMakeLake(METALAKE_NAME);
+ PolicyMetaService policyMetaService = PolicyMetaService.getInstance();
+ PolicyEntity policy =
+ createPolicy(
+ RandomIdGenerator.INSTANCE.nextId(),
+ NamespaceUtil.ofPolicy(METALAKE_NAME),
+ "hook-delete-rollback",
+ AUDIT_INFO);
+ policyMetaService.insertPolicy(policy, false);
+
+ PolicyLifecycleHook hook =
+ new PolicyLifecycleHook() {
+ @Override
+ public void beforeRemove(long policyId) {
+ throw new IllegalArgumentException("bad cron on delete");
+ }
+ };
+ PolicyLifecycleHooks.register(hook);
+ try {
+ assertThrows(
+ IllegalArgumentException.class,
+ () -> policyMetaService.deletePolicy(policy.nameIdentifier()));
+ assertTrue(backend.exists(policy.nameIdentifier(),
Entity.EntityType.POLICY));
+ } finally {
+ PolicyLifecycleHooks.unregister(hook);
+ }
+ }
+
+ @TestTemplate
+ public void testOverwriteFiresUpdateHooksWithRetainedPolicyId() throws
IOException {
+ createAndInsertMakeLake(METALAKE_NAME);
+ PolicyMetaService policyMetaService = PolicyMetaService.getInstance();
+ PolicyEntity existing =
+ createPolicy(
+ RandomIdGenerator.INSTANCE.nextId(),
+ NamespaceUtil.ofPolicy(METALAKE_NAME),
+ "hook-overwrite-id",
+ AUDIT_INFO);
+ policyMetaService.insertPolicy(existing, false);
+
+ List<String> events = new ArrayList<>();
+ PolicyLifecycleHook hook =
+ new PolicyLifecycleHook() {
+ @Override
+ public void beforePersist(PolicyEntity policy) {
+ events.add("beforePersist " + policy.id());
+ }
+
+ @Override
+ public void afterPersist(PolicyEntity policy) {
+ events.add("afterPersist " + policy.id());
+ }
+
+ @Override
+ public void beforeUpdate(PolicyEntity policy) {
+ events.add("beforeUpdate " + policy.id());
+ }
+
+ @Override
+ public void afterUpdate(PolicyEntity policy) {
+ events.add("afterUpdate " + policy.id());
+ }
+ };
+ PolicyLifecycleHooks.register(hook);
+ try {
+ long incomingId = RandomIdGenerator.INSTANCE.nextId();
+ PolicyEntity sameNameOtherId =
+ createPolicy(
+ incomingId, NamespaceUtil.ofPolicy(METALAKE_NAME),
"hook-overwrite-id", AUDIT_INFO);
+ policyMetaService.insertPolicy(sameNameOtherId, true);
+
+ assertEquals(
+ List.of("beforeUpdate " + existing.id(), "afterUpdate " +
existing.id()), events);
+ assertFalse(events.stream().anyMatch(e ->
e.contains(String.valueOf(incomingId))));
+ } finally {
+ PolicyLifecycleHooks.unregister(hook);
+ }
+ }
+
@TestTemplate
public void testInsertAlreadyExistsException() throws IOException {
createAndInsertMakeLake(METALAKE_NAME);
diff --git
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandLifecycleHook.java
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandLifecycleHook.java
new file mode 100644
index 0000000000..e38c46bef9
--- /dev/null
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandLifecycleHook.java
@@ -0,0 +1,203 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.policy;
+
+import com.github.kagkarlsson.scheduler.task.ExecutionComplete;
+import com.github.kagkarlsson.scheduler.task.ExecutionOperations;
+import java.time.Instant;
+import java.util.Optional;
+import javax.annotation.Nullable;
+import org.apache.gravitino.meta.BaseMetalake;
+import org.apache.gravitino.meta.PolicyEntity;
+import org.apache.gravitino.policy.ExecutablePolicyContent;
+import org.apache.gravitino.policy.Policy;
+import org.apache.gravitino.policy.PolicyContent;
+import org.apache.gravitino.policy.hook.PolicyLifecycleHook;
+import org.apache.gravitino.storage.relational.service.MetalakeMetaService;
+import org.apache.gravitino.storage.relational.service.PolicyMetaService;
+import org.apache.gravitino.utils.NamespaceUtil;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Keeps the expand-pool row aligned with a maintenance policy inside the
policy transaction.
+ *
+ * <p>An enabled maintenance policy with a non-blank {@code timeTrigger} is
scheduled. Disable, a
+ * blank trigger, and a non-maintenance write remove the row. Delete removes
the row by policy id.
+ * An invalid cron throws on write so the policy change rolls back; on startup
reconcile the expand
+ * row is removed so scheduler state is corrected while the bad policy content
remains to be edited.
+ */
+public final class PolicyExpandLifecycleHook implements PolicyLifecycleHook {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(PolicyExpandLifecycleHook.class);
+
+ private final PolicyExpandRegistrar registrar;
+
+ /**
+ * @param registrar registers policy-expand task rows on the current SQL
session
+ */
+ public PolicyExpandLifecycleHook(PolicyExpandRegistrar registrar) {
+ this.registrar = registrar;
+ }
+
+ @Override
+ public void afterPersist(PolicyEntity policy) {
+ syncExpandRow(policy);
+ }
+
+ @Override
+ public void afterUpdate(PolicyEntity policy) {
+ syncExpandRow(policy);
+ }
+
+ @Override
+ public void beforeRemove(long policyId) {
+ registrar.remove(policyId);
+ }
+
+ /**
+ * Aligns the {@code policy-expand} task for a stored policy (startup
reconcile).
+ *
+ * <p>Only mutates expand tasks: inserts when missing; reschedules a
<em>future</em> due that no
+ * longer matches the policy cron; leaves an overdue due (< now) alone so
the expand pool can
+ * catch up after downtime. Removes the expand task when the policy must not
run or the cron is
+ * invalid. Never deletes the policy itself.
+ *
+ * @param policy persisted policy
+ */
+ public void align(PolicyEntity policy) {
+ if (!isMaintenance(policy.content())) {
+ registrar.remove(policy.id());
+ return;
+ }
+ String trigger = timeTrigger(policy.content());
+ if (!policy.enabled() || trigger == null || trigger.isBlank()) {
+ registrar.remove(policy.id());
+ return;
+ }
+ Instant expectedNext;
+ try {
+ expectedNext = PolicyExpandRegistrar.nextExecution(trigger);
+ } catch (IllegalArgumentException e) {
+ registrar.remove(policy.id());
+ LOG.error(
+ "Cleared policy-expand task for policy id {} name {} with invalid
timeTrigger; fix the policy content",
+ policy.id(),
+ policy.name(),
+ e);
+ return;
+ }
+ Optional<Instant> current = registrar.scheduledExecutionTime(policy.id());
+ Instant now = Instant.now();
+ if (current.isPresent() && current.get().isBefore(now)) {
+ // Overdue: keep so db-scheduler can catch up after a restart.
+ return;
+ }
+ if (current.isPresent() && current.get().equals(expectedNext)) {
+ return;
+ }
+ registrar.scheduleAt(policy.id(), expectedNext);
+ }
+
+ /**
+ * Finishes one pick of {@code policy-expand}.
+ *
+ * <p>Reloads the policy and either removes the row or moves it to the next
cron fire. An empty
+ * completion would leave the row picked, and db-scheduler then rejects
alter and delete of that
+ * policy.
+ *
+ * @param instanceId db-scheduler task instance id, the policy id
+ * @param complete the pick that just finished
+ * @param operations completion actions for that pick
+ */
+ public static void completePick(
+ String instanceId, ExecutionComplete complete, ExecutionOperations<Void>
operations) {
+ long policyId;
+ try {
+ policyId = Long.parseLong(instanceId);
+ } catch (NumberFormatException ignored) {
+ operations.remove();
+ return;
+ }
+ PolicyEntity policy = findPolicy(policyId);
+ if (policy == null || !isMaintenance(policy.content())) {
+ operations.remove();
+ return;
+ }
+ String trigger = timeTrigger(policy.content());
+ if (!policy.enabled() || trigger == null || trigger.isBlank()) {
+ operations.remove();
+ return;
+ }
+ try {
+ operations.reschedule(complete,
PolicyExpandRegistrar.nextExecution(trigger));
+ } catch (IllegalArgumentException e) {
+ LOG.error(
+ "Removing policy-expand {} after invalid timeTrigger on policy {}",
+ policyId,
+ policy.name(),
+ e);
+ operations.remove();
+ }
+ }
+
+ private void syncExpandRow(PolicyEntity policy) {
+ if (!isMaintenance(policy.content())) {
+ registrar.remove(policy.id());
+ return;
+ }
+ registrar.sync(policy.id(), policy.enabled(),
timeTrigger(policy.content()));
+ }
+
+ private static boolean isMaintenance(@Nullable PolicyContent content) {
+ if (content == null) {
+ return false;
+ }
+ for (Policy.BuiltInType type : Policy.BuiltInType.values()) {
+ if (MaintenancePolicyDefinitions.isMaintenance(type)
+ && type.contentClass().isInstance(content)) {
+ return true;
+ }
+ }
+ return false;
+ }
+
+ @Nullable
+ private static PolicyEntity findPolicy(long policyId) {
+ MetalakeMetaService metalakes = MetalakeMetaService.getInstance();
+ PolicyMetaService policies = PolicyMetaService.getInstance();
+ for (BaseMetalake metalake : metalakes.listMetalakes()) {
+ for (PolicyEntity policy :
+
policies.listPoliciesByNamespace(NamespaceUtil.ofPolicy(metalake.name()))) {
+ if (policy.id() == policyId) {
+ return policy;
+ }
+ }
+ }
+ return null;
+ }
+
+ @Nullable
+ private static String timeTrigger(@Nullable PolicyContent content) {
+ if (!(content instanceof ExecutablePolicyContent)) {
+ return null;
+ }
+ return ((ExecutablePolicyContent) content).executionSpec().timeTrigger();
+ }
+}
diff --git
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandReconciler.java
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandReconciler.java
new file mode 100644
index 0000000000..21b5e68b58
--- /dev/null
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandReconciler.java
@@ -0,0 +1,118 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+ * KIND, either express or implied. See the License for the
+ * specific language governing permissions and limitations
+ * under the License.
+ */
+package org.apache.gravitino.maintenance.policy;
+
+import java.util.HashSet;
+import java.util.Set;
+import org.apache.gravitino.meta.BaseMetalake;
+import org.apache.gravitino.meta.PolicyEntity;
+import org.apache.gravitino.storage.relational.service.MetalakeMetaService;
+import org.apache.gravitino.storage.relational.service.PolicyMetaService;
+import org.apache.gravitino.storage.relational.utils.SessionUtils;
+import org.apache.gravitino.utils.NamespaceUtil;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Reconciles {@code policy-expand} tasks when TMS starts.
+ *
+ * <p>Only {@code scheduled_expand_tasks} rows are touched — never policy rows
or other scheduler
+ * tables. Lifecycle hooks cover later writes; startup must still converge
existing state:
+ *
+ * <ul>
+ * <li><b>missing</b> — enabled maintenance policy with a valid {@code
timeTrigger} but no expand
+ * task → insert the next due after now
+ * <li><b>wrong future due</b> — expand due is still in the future but does
not match the policy
+ * cron → reschedule; invalid cron → remove that expand task only
+ * <li><b>overdue</b> — expand due is before now → leave unchanged so the
pool can catch up
+ * <li><b>surplus</b> — expand task for disabled / blank / non-maintenance /
deleted policy →
+ * delete that expand task only
+ * </ul>
+ */
+public final class PolicyExpandReconciler {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(PolicyExpandReconciler.class);
+
+ private final PolicyExpandLifecycleHook aligner;
+ private final PolicyExpandRegistrar registrar;
+ private final MetalakeMetaService metalakeMetaService;
+ private final PolicyMetaService policyMetaService;
+
+ /**
+ * @param aligner same expand rules as the lifecycle hook
+ * @param registrar writer used for orphan cleanup; both must use {@link
+ * PolicyExpandRegistrar#joiningCurrentTransaction()}
+ */
+ public PolicyExpandReconciler(
+ PolicyExpandLifecycleHook aligner, PolicyExpandRegistrar registrar) {
+ this(aligner, registrar, MetalakeMetaService.getInstance(),
PolicyMetaService.getInstance());
+ }
+
+ /**
+ * @param aligner expand aligner
+ * @param registrar expand registrar for orphan deletes
+ * @param metalakeMetaService metalake listing
+ * @param policyMetaService policy listing
+ */
+ public PolicyExpandReconciler(
+ PolicyExpandLifecycleHook aligner,
+ PolicyExpandRegistrar registrar,
+ MetalakeMetaService metalakeMetaService,
+ PolicyMetaService policyMetaService) {
+ this.aligner = aligner;
+ this.registrar = registrar;
+ this.metalakeMetaService = metalakeMetaService;
+ this.policyMetaService = policyMetaService;
+ }
+
+ /**
+ * For each policy, insert or correct its expand task; then delete orphan
expand tasks.
+ *
+ * @return number of policies walked
+ */
+ public int reconcileAll() {
+ Set<Long> knownPolicyIds = new HashSet<>();
+ int aligned = 0;
+ // Missing → insert; wrong due / invalid cron → correct or clear expand
task only.
+ for (BaseMetalake metalake : metalakeMetaService.listMetalakes()) {
+ for (PolicyEntity policy :
+
policyMetaService.listPoliciesByNamespace(NamespaceUtil.ofPolicy(metalake.name())))
{
+ knownPolicyIds.add(policy.id());
+ SessionUtils.doMultipleWithCommit(() -> aligner.align(policy));
+ aligned++;
+ }
+ }
+ // Surplus expand tasks with no policy left → delete expand only.
+ int[] removedOrphans = {0};
+ SessionUtils.doMultipleWithCommit(
+ () -> {
+ for (long scheduledId : registrar.scheduledPolicyIds()) {
+ if (!knownPolicyIds.contains(scheduledId)) {
+ registrar.remove(scheduledId);
+ removedOrphans[0]++;
+ }
+ }
+ });
+ LOG.info(
+ "Reconciled expand tasks for {} policies; removed {} orphan expand
tasks",
+ aligned,
+ removedOrphans[0]);
+ return aligned;
+ }
+}
diff --git
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandRegistrar.java
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandRegistrar.java
new file mode 100644
index 0000000000..80b766fd37
--- /dev/null
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/policy/PolicyExpandRegistrar.java
@@ -0,0 +1,325 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.policy;
+
+import com.github.kagkarlsson.scheduler.SchedulerClient;
+import
com.github.kagkarlsson.scheduler.exceptions.TaskInstanceCurrentlyExecutingException;
+import com.github.kagkarlsson.scheduler.task.ExecutionComplete;
+import com.github.kagkarlsson.scheduler.task.TaskInstance;
+import com.github.kagkarlsson.scheduler.task.TaskInstanceId;
+import com.github.kagkarlsson.scheduler.task.schedule.CronSchedule;
+import com.github.kagkarlsson.scheduler.task.schedule.CronStyle;
+import com.github.kagkarlsson.scheduler.task.schedule.Schedule;
+import com.google.common.annotations.VisibleForTesting;
+import java.io.PrintWriter;
+import java.lang.reflect.InvocationHandler;
+import java.lang.reflect.InvocationTargetException;
+import java.lang.reflect.Proxy;
+import java.sql.Connection;
+import java.sql.SQLException;
+import java.sql.SQLFeatureNotSupportedException;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.util.HashSet;
+import java.util.Optional;
+import java.util.Set;
+import javax.annotation.Nullable;
+import javax.sql.DataSource;
+import
org.apache.gravitino.maintenance.scheduler.TableMaintenanceSchedulerPools;
+import org.apache.gravitino.storage.relational.session.SqlSessions;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+/**
+ * Registers {@code policy-expand} tasks in {@link
+ * TableMaintenanceSchedulerPools#EXPAND_TASKS_TABLE} through a db-scheduler
client.
+ *
+ * <p>The instance id is the policy id. A blank {@code timeTrigger}, or a
disabled policy, removes
+ * the row. Handlers that fan the row out into cron jobs are out of scope here.
+ *
+ * <p>{@code timeTrigger} is a six-field Quartz cron. This class decides
whether the expression has
+ * a next fire time; policy content does not parse cron itself.
+ *
+ * <p>{@link #joiningCurrentTransaction()} writes through the open MyBatis
session. The client must
+ * not commit or close that connection, so the policy transaction and the
expand row commit
+ * together.
+ */
+public final class PolicyExpandRegistrar {
+
+ private static final Logger LOG =
LoggerFactory.getLogger(PolicyExpandRegistrar.class);
+
+ /** Attempts while a pick still owns the expand row. */
+ private static final int EXECUTING_ATTEMPTS = 10;
+
+ /** Pause between attempts, so a short pick can finish inside the policy
transaction. */
+ private static final long EXECUTING_WAIT_MS = 50L;
+
+ private final DataSource dataSource;
+
+ private PolicyExpandRegistrar(DataSource dataSource) {
+ this.dataSource = dataSource;
+ }
+
+ /**
+ * Opens a client that writes on the current thread's open SQL session.
+ *
+ * <p>The db-scheduler client is created on each call, because building it
borrows a connection.
+ * That connection is the open session, so this registrar must be used
inside a policy
+ * transaction.
+ *
+ * @return registrar that inserts, moves, and deletes expand rows in the
caller's transaction
+ */
+ public static PolicyExpandRegistrar joiningCurrentTransaction() {
+ return new PolicyExpandRegistrar(new SessionDataSource());
+ }
+
+ /**
+ * Opens a client that reads committed rows from the entity-store pool.
+ *
+ * @param dataSource relational JDBC pool shared with the entity store
+ * @return registrar that can read expand rows after the policy transaction
commits
+ */
+ public static PolicyExpandRegistrar forReads(DataSource dataSource) {
+ return new PolicyExpandRegistrar(dataSource);
+ }
+
+ private SchedulerClient client() {
+ return SchedulerClient.Builder.create(
+ dataSource, TableMaintenanceSchedulerPools.policyExpandTask())
+ .tableName(TableMaintenanceSchedulerPools.EXPAND_TASKS_TABLE)
+ .build();
+ }
+
+ /**
+ * Inserts or moves the expand row to the next cron fire time, or removes it
when the policy
+ * should not run.
+ *
+ * @param policyId persisted policy id, used as the db-scheduler task
instance
+ * @param enabled whether the policy is enabled
+ * @param timeTrigger cron expression, or {@code null} when the policy has
no time trigger
+ */
+ public void sync(long policyId, boolean enabled, @Nullable String
timeTrigger) {
+ TaskInstanceId instanceId = instanceId(policyId);
+ if (!enabled || timeTrigger == null || timeTrigger.isBlank()) {
+ cancelIfPresent(instanceId);
+ return;
+ }
+ scheduleAt(policyId, nextExecution(timeTrigger));
+ }
+
+ /**
+ * Inserts or moves the expand row to an explicit due time.
+ *
+ * @param policyId persisted policy id
+ * @param due execution time to store
+ */
+ @VisibleForTesting
+ void scheduleAt(long policyId, Instant due) {
+ runDespiteExecution(() -> scheduleOnce(policyId, due));
+ }
+
+ /**
+ * Deletes the expand row for a dropped policy.
+ *
+ * @param policyId persisted policy id
+ */
+ public void remove(long policyId) {
+ cancelIfPresent(instanceId(policyId));
+ }
+
+ /**
+ * @param policyId persisted policy id
+ * @return next execution time of the expand row, or empty when the row is
absent
+ */
+ Optional<Instant> scheduledExecutionTime(long policyId) {
+ return client()
+ .getScheduledExecution(instanceId(policyId))
+ .map(execution -> execution.getExecutionTime());
+ }
+
+ /**
+ * @return policy ids that currently have a {@code policy-expand} row
+ */
+ Set<Long> scheduledPolicyIds() {
+ Set<Long> ids = new HashSet<>();
+ client()
+ .fetchScheduledExecutionsForTask(
+ TableMaintenanceSchedulerPools.POLICY_EXPAND,
+ Void.class,
+ execution -> {
+ try {
+ ids.add(Long.parseLong(execution.getTaskInstance().getId()));
+ } catch (NumberFormatException ignored) {
+ // Non-numeric instance ids are not policy-expand leases from
this registrar.
+ }
+ });
+ return ids;
+ }
+
+ /**
+ * Next fire time for a Quartz cron, strictly after now.
+ *
+ * @param timeTrigger six-field Quartz cron expression
+ * @return next execution instant in the system default zone
+ */
+ static Instant nextExecution(String timeTrigger) {
+ CronSchedule schedule = new CronSchedule(timeTrigger,
ZoneId.systemDefault(), CronStyle.QUARTZ);
+ Instant next =
schedule.getNextExecutionTime(ExecutionComplete.simulatedSuccess(Instant.now()));
+ if (next == null || Schedule.NEVER.equals(next)) {
+ throw new IllegalArgumentException("timeTrigger has no next execution: "
+ timeTrigger);
+ }
+ return next;
+ }
+
+ private static TaskInstanceId instanceId(long policyId) {
+ return TaskInstanceId.of(TableMaintenanceSchedulerPools.POLICY_EXPAND,
Long.toString(policyId));
+ }
+
+ private void scheduleOnce(long policyId, Instant due) {
+ TaskInstanceId instanceId = instanceId(policyId);
+ boolean scheduled =
+ client()
+ .schedule(
+ new TaskInstance<>(
+ TableMaintenanceSchedulerPools.POLICY_EXPAND,
instanceId.getId()),
+ due,
+ SchedulerClient.ScheduleOptions.WHEN_EXISTS_RESCHEDULE);
+ if (!scheduled) {
+ throw new IllegalStateException(
+ "Failed to schedule policy-expand instance " + instanceId.getId() +
" at " + due);
+ }
+ }
+
+ private void cancelIfPresent(TaskInstanceId instanceId) {
+ runDespiteExecution(
+ () -> {
+ SchedulerClient client = client();
+ if (client.getScheduledExecution(instanceId).isPresent()) {
+ client.cancel(instanceId);
+ }
+ });
+ }
+
+ /**
+ * Retries while db-scheduler reports the expand row as picked.
+ *
+ * <p>A pick rejects both reschedule and cancel. Failing the attempt here
would roll the policy
+ * transaction back. After the retry window the in-flight execution applies
the committed policy
+ * from {@link PolicyExpandLifecycleHook#completePick}.
+ *
+ * @param action one scheduler write
+ */
+ private static void runDespiteExecution(Runnable action) {
+ for (int attempt = 1; ; attempt++) {
+ try {
+ action.run();
+ return;
+ } catch (TaskInstanceCurrentlyExecutingException e) {
+ if (attempt >= EXECUTING_ATTEMPTS) {
+ LOG.warn(
+ "policy-expand row is still executing; leaving the in-flight
pick to apply the policy",
+ e);
+ return;
+ }
+ try {
+ Thread.sleep(EXECUTING_WAIT_MS);
+ } catch (InterruptedException interrupted) {
+ Thread.currentThread().interrupt();
+ e.addSuppressed(interrupted);
+ throw e;
+ }
+ }
+ }
+ }
+
+ /**
+ * Hands db-scheduler the connection of the open SQL session. {@code close},
{@code commit}, and
+ * {@code rollback} are ignored so the session owner commits the expand row
with the policy.
+ */
+ private static final class SessionDataSource implements DataSource {
+
+ @Override
+ public Connection getConnection() throws SQLException {
+ return nonClosing(SqlSessions.currentConnection());
+ }
+
+ @Override
+ public Connection getConnection(String username, String password) throws
SQLException {
+ return getConnection();
+ }
+
+ @Override
+ public PrintWriter getLogWriter() {
+ return null;
+ }
+
+ @Override
+ public void setLogWriter(PrintWriter out) {}
+
+ @Override
+ public void setLoginTimeout(int seconds) {}
+
+ @Override
+ public int getLoginTimeout() {
+ return 0;
+ }
+
+ @Override
+ public java.util.logging.Logger getParentLogger() throws
SQLFeatureNotSupportedException {
+ throw new SQLFeatureNotSupportedException();
+ }
+
+ @Override
+ public <T> T unwrap(Class<T> iface) throws SQLException {
+ if (iface.isInstance(this)) {
+ return iface.cast(this);
+ }
+ throw new SQLException("Not a wrapper of " + iface.getName());
+ }
+
+ @Override
+ public boolean isWrapperFor(Class<?> iface) {
+ return iface.isInstance(this);
+ }
+
+ private static Connection nonClosing(Connection connection) {
+ InvocationHandler handler =
+ (proxy, method, args) -> {
+ String name = method.getName();
+ if ("close".equals(name) || "commit".equals(name) ||
"rollback".equals(name)) {
+ return null;
+ }
+ if ("setAutoCommit".equals(name)) {
+ return null;
+ }
+ if ("getAutoCommit".equals(name)) {
+ return false;
+ }
+ try {
+ return method.invoke(connection, args);
+ } catch (InvocationTargetException e) {
+ throw e.getCause();
+ }
+ };
+ return (Connection)
+ Proxy.newProxyInstance(
+ Connection.class.getClassLoader(), new Class<?>[]
{Connection.class}, handler);
+ }
+ }
+}
diff --git
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/TableMaintenanceSchedulerPools.java
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/TableMaintenanceSchedulerPools.java
index e60bbf9d18..ebf608129e 100644
---
a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/TableMaintenanceSchedulerPools.java
+++
b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/TableMaintenanceSchedulerPools.java
@@ -27,6 +27,7 @@ import java.util.ArrayList;
import java.util.List;
import javax.sql.DataSource;
import org.apache.gravitino.maintenance.TableMaintenanceConfig;
+import org.apache.gravitino.maintenance.policy.PolicyExpandLifecycleHook;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -37,8 +38,10 @@ import org.slf4j.LoggerFactory;
* <p>db-scheduler requires a dedicated {@code tableName} per Scheduler
instance for multi-pool
* isolation; sharing one physical table across pools is unsupported (see
design §5.2).
*
- * <p>Task handlers are intentionally minimal in this change; submit and
occupancy logic land in
- * follow-up PRs.
+ * <p>Cron and commit pools poll. {@code policy-expand} rows are written with
a {@link
+ * com.github.kagkarlsson.scheduler.SchedulerClient}; expand polling starts
with fan-out. {@link
+ * #policyExpandTask()} reschedules to the next cron after a pick, so a run
does not leave the row
+ * picked.
*/
public class TableMaintenanceSchedulerPools implements AutoCloseable {
@@ -65,7 +68,11 @@ public class TableMaintenanceSchedulerPools implements
AutoCloseable {
private final List<Scheduler> schedulers = new ArrayList<>();
/**
- * Creates and starts expand, cron, and commit scheduler pools.
+ * Creates and starts the cron and commit scheduler pools.
+ *
+ * <p>The expand pool stays stopped. Due {@code policy-expand} rows would
otherwise be picked by
+ * an empty handler and block policy alter and delete. Expand polling starts
when fan-out can both
+ * run and reschedule the row.
*
* <p>If any pool fails to start, already-started pools are stopped and this
instance returns to
* the idle state so {@link #start} may be retried.
@@ -79,12 +86,6 @@ public class TableMaintenanceSchedulerPools implements
AutoCloseable {
List<Scheduler> built =
List.of(
- buildScheduler(
- dataSource,
- config,
- config.expandThreads(),
- EXPAND_TASKS_TABLE,
- noopTask(POLICY_EXPAND)),
buildScheduler(
dataSource, config, config.cronThreads(), CRON_TASKS_TABLE,
noopTask(CRON_JOB)),
buildScheduler(
@@ -136,10 +137,28 @@ public class TableMaintenanceSchedulerPools implements
AutoCloseable {
.heartbeatInterval(Duration.ofMillis(config.heartbeatIntervalMs()))
.missedHeartbeatsLimit(config.missedHeartbeatsLimit())
.tableName(tableName)
+ .commitWhenAutocommitDisabled(true)
.registerShutdownHook()
.build();
}
+ /**
+ * Expand task definition shared by the registration client and a future
expand pool.
+ *
+ * <p>{@link #start} does not poll this task. When a run does happen,
completion reschedules to
+ * the next cron or removes the row, so the execution does not stay picked.
+ *
+ * @return the {@code policy-expand} task
+ */
+ public static Task<Void> policyExpandTask() {
+ return Tasks.custom(POLICY_EXPAND, Void.class)
+ .execute(
+ (taskInstance, executionContext) ->
+ (executionComplete, executionOperations) ->
+ PolicyExpandLifecycleHook.completePick(
+ taskInstance.getId(), executionComplete,
executionOperations));
+ }
+
private static Task<Void> noopTask(String taskName) {
return Tasks.custom(taskName, Void.class)
.execute(
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 ec9b9511bc..3b3f7320ee 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
@@ -27,7 +27,11 @@ 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.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.policy.hook.PolicyLifecycleHooks;
import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper;
import org.glassfish.jersey.server.monitoring.ApplicationEvent;
import org.glassfish.jersey.server.monitoring.ApplicationEventListener;
@@ -41,11 +45,15 @@ import org.slf4j.LoggerFactory;
* 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}.
+ * 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>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).
+ * <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).
*/
@Provider
public class TableMaintenanceRESTFeature implements Feature {
@@ -58,6 +66,9 @@ public class TableMaintenanceRESTFeature implements Feature {
private static final AtomicReference<TableMaintenanceSchedulerPools>
SCHEDULER_POOLS =
new AtomicReference<>();
+ private static final AtomicReference<PolicyExpandLifecycleHook>
LIFECYCLE_HOOK =
+ new AtomicReference<>();
+
@Override
public boolean configure(FeatureContext context) {
if (SCHEDULER_POOLS.get() != null) {
@@ -80,6 +91,16 @@ public class TableMaintenanceRESTFeature implements Feature {
return true;
}
+ try {
+ registerLifecycleHook();
+ reconcileExistingPolicies();
+ } catch (RuntimeException e) {
+ unregisterLifecycleHook();
+ SCHEDULER_POOLS.set(null);
+ pools.stop();
+ throw e;
+ }
+
// 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());
@@ -87,8 +108,15 @@ public class TableMaintenanceRESTFeature implements Feature
{
return true;
}
- /** Stops scheduler pools. Invoked from Jersey destroy and tests. */
+ /**
+ * Stops scheduler pools. Invoked from Jersey destroy and tests.
+ *
+ * <p>Unregisters the lifecycle hook first so later policy writes skip
expand sync. In-flight
+ * transactions that already passed the hook may still finish. A later
{@link #configure}
+ * reconciles expand rows again from stored policies.
+ */
public static void stopSchedulers() {
+ unregisterLifecycleHook();
TableMaintenanceSchedulerPools pools = SCHEDULER_POOLS.getAndSet(null);
if (pools != null) {
pools.stop();
@@ -96,6 +124,26 @@ public class TableMaintenanceRESTFeature implements Feature
{
}
}
+ private static void registerLifecycleHook() {
+ PolicyExpandLifecycleHook hook =
+ new
PolicyExpandLifecycleHook(PolicyExpandRegistrar.joiningCurrentTransaction());
+ PolicyLifecycleHooks.register(hook);
+ LIFECYCLE_HOOK.set(hook);
+ }
+
+ private static void reconcileExistingPolicies() {
+ PolicyExpandRegistrar registrar =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ PolicyExpandLifecycleHook aligner = new
PolicyExpandLifecycleHook(registrar);
+ new PolicyExpandReconciler(aligner, registrar).reconcileAll();
+ }
+
+ private static void unregisterLifecycleHook() {
+ PolicyExpandLifecycleHook hook = LIFECYCLE_HOOK.getAndSet(null);
+ if (hook != null) {
+ PolicyLifecycleHooks.unregister(hook);
+ }
+ }
+
private static DataSource resolveEntityStoreDataSource() {
Object dataSource =
SqlSessionFactoryHelper.getInstance()
diff --git
a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/policy/TestPolicyExpandLifecycleHook.java
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/policy/TestPolicyExpandLifecycleHook.java
new file mode 100644
index 0000000000..69336594a1
--- /dev/null
+++
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/policy/TestPolicyExpandLifecycleHook.java
@@ -0,0 +1,135 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.policy;
+
+import java.util.Map;
+import java.util.Set;
+import javax.sql.DataSource;
+import org.apache.gravitino.MetadataObject;
+import org.apache.gravitino.meta.AuditInfo;
+import org.apache.gravitino.meta.PolicyEntity;
+import org.apache.gravitino.policy.IcebergDataCompactionContent;
+import org.apache.gravitino.policy.Policy;
+import org.apache.gravitino.policy.PolicyContent;
+import org.apache.gravitino.policy.PolicyContents;
+import org.apache.gravitino.storage.relational.TestJDBCBackend;
+import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper;
+import org.apache.gravitino.storage.relational.utils.SessionUtils;
+import org.apache.gravitino.utils.NamespaceUtil;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.TestTemplate;
+
+class TestPolicyExpandLifecycleHook extends TestJDBCBackend {
+
+ private static final String CRON = "0 0 2 * * ?";
+
+ @TestTemplate
+ void testEnabledMaintenancePolicyIsScheduled() {
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ PolicyExpandRegistrar reader =
PolicyExpandRegistrar.forReads(entityStoreDataSource());
+ PolicyExpandLifecycleHook hook = new PolicyExpandLifecycleHook(writer);
+
+ SessionUtils.doMultipleWithCommit(() -> hook.afterPersist(policy(42L,
true, compaction(CRON))));
+
+ Assertions.assertTrue(reader.scheduledExecutionTime(42L).isPresent());
+ }
+
+ @TestTemplate
+ void testDisabledMaintenancePolicyRemovesTheRow() {
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ PolicyExpandRegistrar reader =
PolicyExpandRegistrar.forReads(entityStoreDataSource());
+ PolicyExpandLifecycleHook hook = new PolicyExpandLifecycleHook(writer);
+
+ SessionUtils.doMultipleWithCommit(() -> hook.afterPersist(policy(7L, true,
compaction(CRON))));
+ Assertions.assertTrue(reader.scheduledExecutionTime(7L).isPresent());
+
+ SessionUtils.doMultipleWithCommit(() -> hook.afterUpdate(policy(7L, false,
compaction(CRON))));
+ Assertions.assertTrue(reader.scheduledExecutionTime(7L).isEmpty());
+ }
+
+ @TestTemplate
+ void testNonMaintenancePolicyRemovesAnyStrayRow() {
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ PolicyExpandRegistrar reader =
PolicyExpandRegistrar.forReads(entityStoreDataSource());
+ PolicyExpandLifecycleHook hook = new PolicyExpandLifecycleHook(writer);
+ PolicyContent custom = PolicyContents.custom(Map.of(),
Set.of(MetadataObject.Type.TABLE), null);
+
+ SessionUtils.doMultipleWithCommit(() -> writer.sync(1L, true, CRON));
+ Assertions.assertTrue(reader.scheduledExecutionTime(1L).isPresent());
+
+ SessionUtils.doMultipleWithCommit(
+ () -> hook.afterPersist(policy(1L, true, custom,
Policy.BuiltInType.CUSTOM)));
+ Assertions.assertTrue(reader.scheduledExecutionTime(1L).isEmpty());
+ }
+
+ @TestTemplate
+ void testDeleteRemovesByPolicyId() {
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ PolicyExpandRegistrar reader =
PolicyExpandRegistrar.forReads(entityStoreDataSource());
+ PolicyExpandLifecycleHook hook = new PolicyExpandLifecycleHook(writer);
+
+ SessionUtils.doMultipleWithCommit(() -> writer.sync(11L, true, CRON));
+ Assertions.assertTrue(reader.scheduledExecutionTime(11L).isPresent());
+
+ SessionUtils.doMultipleWithCommit(() -> hook.beforeRemove(11L));
+ Assertions.assertTrue(reader.scheduledExecutionTime(11L).isEmpty());
+ }
+
+ private static DataSource entityStoreDataSource() {
+ return SqlSessionFactoryHelper.getInstance()
+ .getSqlSessionFactory()
+ .getConfiguration()
+ .getEnvironment()
+ .getDataSource();
+ }
+
+ private static PolicyEntity policy(long id, boolean enabled, PolicyContent
content) {
+ return policy(id, enabled, content, Policy.BuiltInType.ICEBERG_COMPACTION);
+ }
+
+ private static PolicyEntity policy(
+ long id, boolean enabled, PolicyContent content, Policy.BuiltInType
type) {
+ return PolicyEntity.builder()
+ .withId(id)
+ .withName("nightly")
+ .withNamespace(NamespaceUtil.ofPolicy("ml"))
+ .withPolicyType(type)
+ .withContent(content)
+ .withEnabled(enabled)
+ .withAuditInfo(AuditInfo.EMPTY)
+ .build();
+ }
+
+ private static PolicyContent compaction(String timeTrigger) {
+ return PolicyContents.icebergDataCompaction(
+ IcebergDataCompactionContent.DEFAULT_MIN_DATA_FILE_MSE,
+ IcebergDataCompactionContent.DEFAULT_MIN_DELETE_FILE_NUMBER,
+ IcebergDataCompactionContent.DEFAULT_DATA_FILE_MSE_WEIGHT,
+ IcebergDataCompactionContent.DEFAULT_DELETE_FILE_NUMBER_WEIGHT,
+ IcebergDataCompactionContent.DEFAULT_MAX_PARTITION_NUM,
+ IcebergDataCompactionContent.DEFAULT_REWRITE_STRATEGY,
+ IcebergDataCompactionContent.DEFAULT_SORT_ORDER,
+ IcebergDataCompactionContent.DEFAULT_REWRITE_OPTIONS,
+ timeTrigger,
+ null,
+ null,
+ null,
+ null);
+ }
+}
diff --git
a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/policy/TestPolicyExpandReconciler.java
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/policy/TestPolicyExpandReconciler.java
new file mode 100644
index 0000000000..a2931aafe1
--- /dev/null
+++
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/policy/TestPolicyExpandReconciler.java
@@ -0,0 +1,249 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.policy;
+
+import com.github.kagkarlsson.scheduler.task.ExecutionComplete;
+import com.github.kagkarlsson.scheduler.task.ExecutionOperations;
+import java.io.IOException;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import java.util.List;
+import javax.sql.DataSource;
+import org.apache.gravitino.meta.AuditInfo;
+import org.apache.gravitino.meta.BaseMetalake;
+import org.apache.gravitino.meta.PolicyEntity;
+import org.apache.gravitino.policy.IcebergDataCompactionContent;
+import org.apache.gravitino.policy.Policy;
+import org.apache.gravitino.policy.PolicyContent;
+import org.apache.gravitino.policy.PolicyContents;
+import org.apache.gravitino.storage.RandomIdGenerator;
+import org.apache.gravitino.storage.relational.TestJDBCBackend;
+import org.apache.gravitino.storage.relational.service.MetalakeMetaService;
+import org.apache.gravitino.storage.relational.service.PolicyMetaService;
+import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper;
+import org.apache.gravitino.storage.relational.utils.SessionUtils;
+import org.apache.gravitino.utils.NamespaceUtil;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.TestTemplate;
+import org.mockito.Mockito;
+
+class TestPolicyExpandReconciler extends TestJDBCBackend {
+
+ private static final String CRON = "0 0 2 * * ?";
+ private static final String LATER_CRON = "0 30 4 * * ?";
+ private static final String METALAKE = "metalake_for_expand_reconcile";
+
+ @TestTemplate
+ void testReconcileBackfillsExpandRowForExistingPolicy() throws IOException {
+ createAndInsertMakeLake(METALAKE);
+
+ PolicyEntity policy =
+ PolicyEntity.builder()
+ .withId(RandomIdGenerator.INSTANCE.nextId())
+ .withName("nightly")
+ .withNamespace(NamespaceUtil.ofPolicy(METALAKE))
+ .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION)
+ .withContent(compaction(CRON))
+ .withEnabled(true)
+ .withAuditInfo(AuditInfo.EMPTY)
+ .build();
+ PolicyMetaService.getInstance().insertPolicy(policy, false);
+
+ DataSource dataSource = entityStoreDataSource();
+ PolicyExpandRegistrar reader = PolicyExpandRegistrar.forReads(dataSource);
+
Assertions.assertTrue(reader.scheduledExecutionTime(policy.id()).isEmpty());
+
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ PolicyExpandLifecycleHook aligner = new PolicyExpandLifecycleHook(writer);
+ int aligned =
+ new PolicyExpandReconciler(
+ aligner, writer, MetalakeMetaService.getInstance(),
PolicyMetaService.getInstance())
+ .reconcileAll();
+
+ Assertions.assertTrue(aligned >= 1);
+
Assertions.assertTrue(reader.scheduledExecutionTime(policy.id()).isPresent());
+ }
+
+ @TestTemplate
+ void testReconcilePreservesOverdueExpandDue() throws IOException {
+ createAndInsertMakeLake(METALAKE + "_overdue");
+ String metalake = METALAKE + "_overdue";
+ PolicyEntity policy =
+ PolicyEntity.builder()
+ .withId(RandomIdGenerator.INSTANCE.nextId())
+ .withName("overdue")
+ .withNamespace(NamespaceUtil.ofPolicy(metalake))
+ .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION)
+ .withContent(compaction(CRON))
+ .withEnabled(true)
+ .withAuditInfo(AuditInfo.EMPTY)
+ .build();
+ PolicyMetaService.getInstance().insertPolicy(policy, false);
+
+ // db-scheduler / JDBC timestamps are microsecond precision; truncate so
round-trip equals.
+ Instant overdue =
Instant.now().minusSeconds(3600).truncatedTo(ChronoUnit.MICROS);
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ PolicyExpandRegistrar reader =
PolicyExpandRegistrar.forReads(entityStoreDataSource());
+ SessionUtils.doMultipleWithCommit(() -> writer.scheduleAt(policy.id(),
overdue));
+ Assertions.assertEquals(overdue,
reader.scheduledExecutionTime(policy.id()).orElseThrow());
+
+ PolicyExpandLifecycleHook aligner = new PolicyExpandLifecycleHook(writer);
+ new PolicyExpandReconciler(
+ aligner, writer, MetalakeMetaService.getInstance(),
PolicyMetaService.getInstance())
+ .reconcileAll();
+
+ Assertions.assertEquals(overdue,
reader.scheduledExecutionTime(policy.id()).orElseThrow());
+ }
+
+ @TestTemplate
+ void testReconcileCorrectsStaleExpandDueTime() throws IOException {
+ createAndInsertMakeLake(METALAKE + "_stale");
+ String metalake = METALAKE + "_stale";
+ PolicyEntity policy =
+ PolicyEntity.builder()
+ .withId(RandomIdGenerator.INSTANCE.nextId())
+ .withName("stale-due")
+ .withNamespace(NamespaceUtil.ofPolicy(metalake))
+ .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION)
+ .withContent(compaction(CRON))
+ .withEnabled(true)
+ .withAuditInfo(AuditInfo.EMPTY)
+ .build();
+ PolicyMetaService.getInstance().insertPolicy(policy, false);
+
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ PolicyExpandRegistrar reader =
PolicyExpandRegistrar.forReads(entityStoreDataSource());
+ SessionUtils.doMultipleWithCommit(() -> writer.sync(policy.id(), true,
LATER_CRON));
+ Instant staleDue =
reader.scheduledExecutionTime(policy.id()).orElseThrow();
+
+ PolicyExpandLifecycleHook aligner = new PolicyExpandLifecycleHook(writer);
+ new PolicyExpandReconciler(
+ aligner, writer, MetalakeMetaService.getInstance(),
PolicyMetaService.getInstance())
+ .reconcileAll();
+
+ Instant corrected =
reader.scheduledExecutionTime(policy.id()).orElseThrow();
+ Assertions.assertNotEquals(staleDue, corrected);
+ }
+
+ @TestTemplate
+ void testReconcileRemovesExpandRowForInvalidCron() {
+ PolicyEntity bad =
+ PolicyEntity.builder()
+ .withId(7L)
+ .withName("bad")
+ .withNamespace(NamespaceUtil.ofPolicy("ml"))
+ .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION)
+ .withContent(compaction("not a cron"))
+ .withEnabled(true)
+ .withAuditInfo(AuditInfo.EMPTY)
+ .build();
+ BaseMetalake metalake = Mockito.mock(BaseMetalake.class);
+ Mockito.when(metalake.name()).thenReturn("ml");
+
+ MetalakeMetaService metalakes = Mockito.mock(MetalakeMetaService.class);
+ PolicyMetaService policies = Mockito.mock(PolicyMetaService.class);
+ Mockito.when(metalakes.listMetalakes()).thenReturn(List.of(metalake));
+
Mockito.when(policies.listPoliciesByNamespace(NamespaceUtil.ofPolicy("ml")))
+ .thenReturn(List.of(bad));
+
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ SessionUtils.doMultipleWithCommit(() -> writer.sync(7L, true, CRON));
+ PolicyExpandRegistrar reader =
PolicyExpandRegistrar.forReads(entityStoreDataSource());
+ Assertions.assertTrue(reader.scheduledExecutionTime(7L).isPresent());
+
+ PolicyExpandLifecycleHook aligner = new PolicyExpandLifecycleHook(writer);
+ Assertions.assertDoesNotThrow(
+ () -> new PolicyExpandReconciler(aligner, writer, metalakes,
policies).reconcileAll());
+ Assertions.assertTrue(reader.scheduledExecutionTime(7L).isEmpty());
+ }
+
+ @TestTemplate
+ void testReconcileRemovesOrphanExpandRow() {
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ SessionUtils.doMultipleWithCommit(() -> writer.sync(99L, true, CRON));
+ PolicyExpandRegistrar reader =
PolicyExpandRegistrar.forReads(entityStoreDataSource());
+ Assertions.assertTrue(reader.scheduledExecutionTime(99L).isPresent());
+
+ MetalakeMetaService metalakes = Mockito.mock(MetalakeMetaService.class);
+ PolicyMetaService policies = Mockito.mock(PolicyMetaService.class);
+ Mockito.when(metalakes.listMetalakes()).thenReturn(List.of());
+
+ PolicyExpandLifecycleHook aligner = new PolicyExpandLifecycleHook(writer);
+ new PolicyExpandReconciler(aligner, writer, metalakes,
policies).reconcileAll();
+ Assertions.assertTrue(reader.scheduledExecutionTime(99L).isEmpty());
+ }
+
+ @TestTemplate
+ void testCompletePickReschedulesEnabledPolicy() throws IOException {
+ createAndInsertMakeLake(METALAKE + "_pick");
+ String metalake = METALAKE + "_pick";
+ PolicyEntity policy =
+ PolicyEntity.builder()
+ .withId(RandomIdGenerator.INSTANCE.nextId())
+ .withName("picked")
+ .withNamespace(NamespaceUtil.ofPolicy(metalake))
+ .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION)
+ .withContent(compaction(CRON))
+ .withEnabled(true)
+ .withAuditInfo(AuditInfo.EMPTY)
+ .build();
+ PolicyMetaService.getInstance().insertPolicy(policy, false);
+
+ ExecutionOperations<Void> operations =
Mockito.mock(ExecutionOperations.class);
+ ExecutionComplete complete =
ExecutionComplete.simulatedSuccess(Instant.now());
+ PolicyExpandLifecycleHook.completePick(Long.toString(policy.id()),
complete, operations);
+
+ Mockito.verify(operations)
+ .reschedule(Mockito.eq(complete), Mockito.argThat(due ->
due.isAfter(Instant.now())));
+ Mockito.verify(operations, Mockito.never()).remove();
+ }
+
+ @TestTemplate
+ void testCompletePickRemovesMissingPolicy() {
+ ExecutionOperations<Void> operations =
Mockito.mock(ExecutionOperations.class);
+ ExecutionComplete complete =
ExecutionComplete.simulatedSuccess(Instant.now());
+ PolicyExpandLifecycleHook.completePick("424242", complete, operations);
+ Mockito.verify(operations).remove();
+ }
+
+ private static DataSource entityStoreDataSource() {
+ return SqlSessionFactoryHelper.getInstance()
+ .getSqlSessionFactory()
+ .getConfiguration()
+ .getEnvironment()
+ .getDataSource();
+ }
+
+ private static PolicyContent compaction(String timeTrigger) {
+ return PolicyContents.icebergDataCompaction(
+ IcebergDataCompactionContent.DEFAULT_MIN_DATA_FILE_MSE,
+ IcebergDataCompactionContent.DEFAULT_MIN_DELETE_FILE_NUMBER,
+ IcebergDataCompactionContent.DEFAULT_DATA_FILE_MSE_WEIGHT,
+ IcebergDataCompactionContent.DEFAULT_DELETE_FILE_NUMBER_WEIGHT,
+ IcebergDataCompactionContent.DEFAULT_MAX_PARTITION_NUM,
+ IcebergDataCompactionContent.DEFAULT_REWRITE_STRATEGY,
+ IcebergDataCompactionContent.DEFAULT_SORT_ORDER,
+ IcebergDataCompactionContent.DEFAULT_REWRITE_OPTIONS,
+ timeTrigger,
+ null,
+ null,
+ null,
+ null);
+ }
+}
diff --git
a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/policy/TestPolicyExpandRegistrar.java
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/policy/TestPolicyExpandRegistrar.java
new file mode 100644
index 0000000000..04b93e7530
--- /dev/null
+++
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/policy/TestPolicyExpandRegistrar.java
@@ -0,0 +1,135 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements. See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership. The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing,
+ * software distributed under the License is distributed on an
+ * "AS IS" BASIS, WITHOUT WARRANTIES 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.policy;
+
+import com.github.kagkarlsson.scheduler.SchedulerClient;
+import com.github.kagkarlsson.scheduler.task.TaskInstanceId;
+import java.time.Instant;
+import java.time.temporal.ChronoUnit;
+import javax.sql.DataSource;
+import org.apache.gravitino.maintenance.TableMaintenanceConfig;
+import
org.apache.gravitino.maintenance.scheduler.TableMaintenanceSchedulerPools;
+import org.apache.gravitino.storage.relational.TestJDBCBackend;
+import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper;
+import org.apache.gravitino.storage.relational.utils.SessionUtils;
+import org.junit.jupiter.api.Assertions;
+import org.junit.jupiter.api.TestTemplate;
+
+class TestPolicyExpandRegistrar extends TestJDBCBackend {
+
+ private static final String CRON = "0 0 2 * * ?";
+ private static final String LATER_CRON = "0 30 4 * * ?";
+
+ @TestTemplate
+ void testSyncInsertsReschedulesAndRemovesTheExpandRow() {
+ DataSource dataSource = entityStoreDataSource();
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ PolicyExpandRegistrar reader = PolicyExpandRegistrar.forReads(dataSource);
+
+ SessionUtils.doMultipleWithCommit(() -> writer.sync(42L, true, CRON));
+ Instant firstDue = executionTime(reader, 42L);
+ Assertions.assertTrue(firstDue.isAfter(Instant.now()));
+
+ SessionUtils.doMultipleWithCommit(() -> writer.sync(42L, true,
LATER_CRON));
+ Instant secondDue = executionTime(reader, 42L);
+ Assertions.assertNotEquals(firstDue, secondDue);
+
+ SessionUtils.doMultipleWithCommit(() -> writer.sync(42L, false,
LATER_CRON));
+ Assertions.assertTrue(reader.scheduledExecutionTime(42L).isEmpty());
+
+ SessionUtils.doMultipleWithCommit(() -> writer.sync(42L, true, CRON));
+ Assertions.assertTrue(reader.scheduledExecutionTime(42L).isPresent());
+ SessionUtils.doMultipleWithCommit(() -> writer.remove(42L));
+ Assertions.assertTrue(reader.scheduledExecutionTime(42L).isEmpty());
+ SessionUtils.doMultipleWithCommit(() -> writer.remove(42L));
+ }
+
+ @TestTemplate
+ void testFailedTransactionRollsBackTheExpandRow() {
+ DataSource dataSource = entityStoreDataSource();
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ PolicyExpandRegistrar reader = PolicyExpandRegistrar.forReads(dataSource);
+
+ Assertions.assertThrows(
+ IllegalStateException.class,
+ () ->
+ SessionUtils.doMultipleWithCommit(
+ () -> writer.sync(99L, true, CRON),
+ () -> {
+ throw new IllegalStateException("boom");
+ }));
+
+ Assertions.assertTrue(reader.scheduledExecutionTime(99L).isEmpty());
+ }
+
+ @TestTemplate
+ void testInvalidCronFailsBeforeAnExpandRowIsWritten() {
+ DataSource dataSource = entityStoreDataSource();
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ PolicyExpandRegistrar reader = PolicyExpandRegistrar.forReads(dataSource);
+
+ Assertions.assertThrows(
+ IllegalArgumentException.class,
+ () -> SessionUtils.doMultipleWithCommit(() -> writer.sync(100L, true,
"not a cron")));
+
+ Assertions.assertTrue(reader.scheduledExecutionTime(100L).isEmpty());
+ }
+
+ @TestTemplate
+ void testRunningPoolsLeaveOverdueExpandRowsUnpicked() throws Exception {
+ Instant overdue =
Instant.now().minusSeconds(30).truncatedTo(ChronoUnit.MICROS);
+ PolicyExpandRegistrar writer =
PolicyExpandRegistrar.joiningCurrentTransaction();
+ DataSource dataSource = entityStoreDataSource();
+ SessionUtils.doMultipleWithCommit(() -> writer.scheduleAt(11L, overdue));
+
+ SchedulerClient client =
+ SchedulerClient.Builder.create(
+ dataSource, TableMaintenanceSchedulerPools.policyExpandTask())
+ .tableName(TableMaintenanceSchedulerPools.EXPAND_TASKS_TABLE)
+ .build();
+ TaskInstanceId instanceId =
+ TaskInstanceId.of(TableMaintenanceSchedulerPools.POLICY_EXPAND, "11");
+ Instant stored =
client.getScheduledExecution(instanceId).orElseThrow().getExecutionTime();
+
+ TableMaintenanceSchedulerPools pools = new
TableMaintenanceSchedulerPools();
+ pools.start(dataSource, TableMaintenanceConfig.forTest(1, 1, 1, 100L,
1_000L, 4, 10));
+ try {
+ Thread.sleep(400L);
+
Assertions.assertFalse(client.getScheduledExecution(instanceId).orElseThrow().isPicked());
+ Assertions.assertEquals(
+ stored,
client.getScheduledExecution(instanceId).orElseThrow().getExecutionTime());
+ } finally {
+ pools.stop();
+ }
+ }
+
+ private static DataSource entityStoreDataSource() {
+ return SqlSessionFactoryHelper.getInstance()
+ .getSqlSessionFactory()
+ .getConfiguration()
+ .getEnvironment()
+ .getDataSource();
+ }
+
+ private static Instant executionTime(PolicyExpandRegistrar registrar, long
policyId) {
+ return registrar
+ .scheduledExecutionTime(policyId)
+ .orElseThrow(() -> new AssertionError("missing expand row " +
policyId));
+ }
+}
diff --git
a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestTableMaintenanceSchedulerPools.java
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestTableMaintenanceSchedulerPools.java
index 5de4f541b1..f654612bbb 100644
---
a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestTableMaintenanceSchedulerPools.java
+++
b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestTableMaintenanceSchedulerPools.java
@@ -39,7 +39,7 @@ class TestTableMaintenanceSchedulerPools extends
TestJDBCBackend {
}
@TestTemplate
- void testThreePoolsStartAndStop() {
+ void testCronAndCommitPoolsStartAndStop() {
DataSource dataSource =
SqlSessionFactoryHelper.getInstance()
.getSqlSessionFactory()
