This is an automated email from the ASF dual-hosted git repository. roryqi pushed a commit to branch main in repository https://gitbox.apache.org/repos/asf/gravitino.git
commit 43e3933c8b94667b24d97225a39b7c0e5914ef63 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()
