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 c4ab2879a70bc7f284a8598be6a456d32b88b5c7 Author: remilu <[email protected]> AuthorDate: Mon Oct 26 07:55:09 2026 +0800 [#14217] feat(tms): Relational cron-job enqueue with snapshot batches Fixes #14217. Squashed from fork #76 (B6). --- .../scheduler/TableMaintenanceSchedulerPools.java | 41 +- .../maintenance/scheduler/job/cron/CronJobIds.java | 206 ++++++++++ .../scheduler/job/cron/CronJobScheduler.java | 5 +- .../job/cron/PendingCronExpandOccupancyGate.java | 51 +++ .../job/cron/RelationalCronJobScheduler.java | 427 +++++++++++++++++++++ .../scheduler/job/cron/TestCronJobIds.java | 58 +++ .../scheduler/job/cron/TestCronJobScheduler.java | 325 ++++++++++++++++ 7 files changed, 1090 insertions(+), 23 deletions(-) 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 ebf608129e..88582781e1 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 @@ -19,6 +19,7 @@ package org.apache.gravitino.maintenance.scheduler; import com.github.kagkarlsson.scheduler.Scheduler; +import com.github.kagkarlsson.scheduler.serializer.JacksonSerializer; import com.github.kagkarlsson.scheduler.task.Task; import com.github.kagkarlsson.scheduler.task.helper.Tasks; import com.google.common.base.Preconditions; @@ -26,8 +27,8 @@ import java.time.Duration; import java.util.ArrayList; import java.util.List; import javax.sql.DataSource; +import org.apache.gravitino.json.JsonUtils; import org.apache.gravitino.maintenance.TableMaintenanceConfig; -import org.apache.gravitino.maintenance.policy.PolicyExpandLifecycleHook; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -38,10 +39,8 @@ 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>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. + * <p>Cron enqueue uses {@link #cronJobTask()} with Jackson {@code task_data}; pick submit lands in + * a follow-up. */ public class TableMaintenanceSchedulerPools implements AutoCloseable { @@ -68,11 +67,7 @@ public class TableMaintenanceSchedulerPools implements AutoCloseable { private final List<Scheduler> schedulers = new ArrayList<>(); /** - * 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. + * Creates and starts expand, cron, and commit scheduler pools. * * <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. @@ -87,7 +82,13 @@ public class TableMaintenanceSchedulerPools implements AutoCloseable { List<Scheduler> built = List.of( buildScheduler( - dataSource, config, config.cronThreads(), CRON_TASKS_TABLE, noopTask(CRON_JOB)), + dataSource, + config, + config.expandThreads(), + EXPAND_TASKS_TABLE, + noopTask(POLICY_EXPAND)), + buildScheduler( + dataSource, config, config.cronThreads(), CRON_TASKS_TABLE, cronJobTask()), buildScheduler( dataSource, config, @@ -130,12 +131,13 @@ public class TableMaintenanceSchedulerPools implements AutoCloseable { TableMaintenanceConfig config, int threads, String tableName, - Task<Void> task) { + Task<?> task) { return Scheduler.create(dataSource, task) .threads(threads) .pollingInterval(Duration.ofMillis(config.pollingIntervalMs())) .heartbeatInterval(Duration.ofMillis(config.heartbeatIntervalMs())) .missedHeartbeatsLimit(config.missedHeartbeatsLimit()) + .serializer(new JacksonSerializer(JsonUtils.anyFieldMapper())) .tableName(tableName) .commitWhenAutocommitDisabled(true) .registerShutdownHook() @@ -143,20 +145,15 @@ public class TableMaintenanceSchedulerPools implements AutoCloseable { } /** - * 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. + * Cron-pool task stub: removes the one-shot row after dequeue. Real submit lands in a follow-up. * - * @return the {@code policy-expand} task + * @return the {@code cron-job} task */ - public static Task<Void> policyExpandTask() { - return Tasks.custom(POLICY_EXPAND, Void.class) + public static Task<MaintenanceJobSpec> cronJobTask() { + return Tasks.custom(CRON_JOB, MaintenanceJobSpec.class) .execute( (taskInstance, executionContext) -> - (executionComplete, executionOperations) -> - PolicyExpandLifecycleHook.completePick( - taskInstance.getId(), executionComplete, executionOperations)); + (executionComplete, executionOperations) -> executionOperations.remove()); } private static Task<Void> noopTask(String taskName) { diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/CronJobIds.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/CronJobIds.java new file mode 100644 index 0000000000..bc976e0920 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/CronJobIds.java @@ -0,0 +1,206 @@ +/* + * 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.scheduler.job.cron; + +import java.util.Objects; +import java.util.Optional; +import javax.annotation.Nullable; + +/** + * Builds and parses db-scheduler instance ids for one-shot {@code cron-job} rows (§5.2.2). + * + * <p>Most types use {@code table:{table_id}:{policy_id}}. Batched snapshot-expiry uses {@code + * batch:{batch_id}:{policy_id}} where {@code batch_id} is a per-policy sequence ({@code 0}, {@code + * 1}, …) so concurrent batch names stay distinct under one {@code policy_id}. + */ +public final class CronJobIds { + + private static final String TABLE_PREFIX = "table:"; + private static final String BATCH_PREFIX = "batch:"; + + private CronJobIds() {} + + /** + * @param tableId table entity id + * @param policyId policy entity id + * @return cron-job instance id + */ + public static String forTable(long tableId, long policyId) { + return TABLE_PREFIX + tableId + ":" + policyId; + } + + /** + * @param batchId per-policy batch sequence id ({@code 0}, {@code 1}, …) + * @param policyId policy entity id + * @return cron-job instance id {@code batch:{batch_id}:{policy_id}} + */ + public static String forBatch(long batchId, long policyId) { + return BATCH_PREFIX + batchId + ":" + policyId; + } + + /** + * Parses a single-table cron-job instance id. + * + * @param instanceId db-scheduler task instance id + * @return table and policy ids when the id matches {@code table:{tableId}:{policyId}} + */ + public static Optional<TablePolicyIds> parseTable(String instanceId) { + if (instanceId == null || !instanceId.startsWith(TABLE_PREFIX)) { + return Optional.empty(); + } + String rest = instanceId.substring(TABLE_PREFIX.length()); + int sep = rest.indexOf(':'); + if (sep <= 0 || sep == rest.length() - 1) { + return Optional.empty(); + } + try { + long tableId = Long.parseLong(rest.substring(0, sep)); + long policyId = Long.parseLong(rest.substring(sep + 1)); + return Optional.of(new TablePolicyIds(tableId, policyId)); + } catch (NumberFormatException e) { + return Optional.empty(); + } + } + + /** + * Parses a batched snapshot-expiry cron-job instance id. + * + * @param instanceId db-scheduler task instance id + * @return batch and policy ids when the id matches {@code batch:{batchId}:{policyId}} + */ + public static Optional<BatchPolicyIds> parseBatch(String instanceId) { + if (instanceId == null || !instanceId.startsWith(BATCH_PREFIX)) { + return Optional.empty(); + } + String rest = instanceId.substring(BATCH_PREFIX.length()); + int sep = rest.indexOf(':'); + if (sep <= 0 || sep == rest.length() - 1) { + return Optional.empty(); + } + try { + long batchId = Long.parseLong(rest.substring(0, sep)); + long policyId = Long.parseLong(rest.substring(sep + 1)); + return Optional.of(new BatchPolicyIds(batchId, policyId)); + } catch (NumberFormatException e) { + return Optional.empty(); + } + } + + /** Parsed {@code table:{tableId}:{policyId}} ids. */ + public static final class TablePolicyIds { + private final long tableId; + private final long policyId; + + /** + * @param tableId table entity id + * @param policyId policy entity id + */ + public TablePolicyIds(long tableId, long policyId) { + this.tableId = tableId; + this.policyId = policyId; + } + + /** + * @return table entity id + */ + public long tableId() { + return tableId; + } + + /** + * @return policy entity id + */ + public long policyId() { + return policyId; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof TablePolicyIds)) { + return false; + } + TablePolicyIds that = (TablePolicyIds) o; + return tableId == that.tableId && policyId == that.policyId; + } + + @Override + public int hashCode() { + return Objects.hash(tableId, policyId); + } + + @Override + public String toString() { + return forTable(tableId, policyId); + } + } + + /** Parsed {@code batch:{batchId}:{policyId}} ids. */ + public static final class BatchPolicyIds { + private final long batchId; + private final long policyId; + + /** + * @param batchId per-policy batch sequence id + * @param policyId policy entity id + */ + public BatchPolicyIds(long batchId, long policyId) { + this.batchId = batchId; + this.policyId = policyId; + } + + /** + * @return batch sequence id + */ + public long batchId() { + return batchId; + } + + /** + * @return policy entity id + */ + public long policyId() { + return policyId; + } + + @Override + public boolean equals(@Nullable Object o) { + if (this == o) { + return true; + } + if (!(o instanceof BatchPolicyIds)) { + return false; + } + BatchPolicyIds that = (BatchPolicyIds) o; + return batchId == that.batchId && policyId == that.policyId; + } + + @Override + public int hashCode() { + return Objects.hash(batchId, policyId); + } + + @Override + public String toString() { + return forBatch(batchId, policyId); + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/CronJobScheduler.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/CronJobScheduler.java index f16cc3febd..dfd6b2ca12 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/CronJobScheduler.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/CronJobScheduler.java @@ -30,6 +30,9 @@ import org.apache.gravitino.meta.PolicyEntity; * batch:{batch_id}:{policy_id}} rows; incomplete batches are filled via {@code * SchedulerClient#getScheduledExecution} / {@code reschedule} before a new {@code batch_id} is * allocated (§5.2.2). + * + * <p>Enqueue builds a {@link org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec} once + * (policy → jobConf) into {@code task_data}; pick consumes that spec. */ @FunctionalInterface public interface CronJobScheduler { @@ -37,7 +40,7 @@ public interface CronJobScheduler { /** * Enqueues due-now cron jobs for the given expand candidates. * - * @param policy effective maintenance policy + * @param policy effective maintenance policy (used to build the jobConf spec) * @param candidates enqueueable tables with winner levels * @param batchCapacity max tables per snapshot-expiry batch row (and max candidates considered) * @return number of newly inserted cron-job rows (merges into existing batches do not count) diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/PendingCronExpandOccupancyGate.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/PendingCronExpandOccupancyGate.java new file mode 100644 index 0000000000..cb9fd6a65a --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/PendingCronExpandOccupancyGate.java @@ -0,0 +1,51 @@ +/* + * 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.scheduler.job.cron; + +import com.google.common.base.Preconditions; +import javax.sql.DataSource; +import org.apache.gravitino.maintenance.scheduler.policy.expand.ExpandOccupancyGate; + +/** + * Decorates an {@link ExpandOccupancyGate} so tables already covered by a queued {@code cron-job} + * row are skipped and do not consume the expand batch budget. + */ +public final class PendingCronExpandOccupancyGate implements ExpandOccupancyGate { + + private final ExpandOccupancyGate delegate; + private final DataSource dataSource; + + /** + * @param delegate in-flight / minInterval gate + * @param dataSource entity-store JDBC pool used by the cron scheduler table + */ + public PendingCronExpandOccupancyGate(ExpandOccupancyGate delegate, DataSource dataSource) { + this.delegate = Preconditions.checkNotNull(delegate, "delegate"); + this.dataSource = Preconditions.checkNotNull(dataSource, "dataSource"); + } + + @Override + public boolean shouldSkip(long tableId, long policyId, long minIntervalMs, long nowMs) { + if (delegate.shouldSkip(tableId, policyId, minIntervalMs, nowMs)) { + return true; + } + return RelationalCronJobScheduler.isQueuedForPolicy( + RelationalCronJobScheduler.client(dataSource), tableId, policyId); + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/RelationalCronJobScheduler.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/RelationalCronJobScheduler.java new file mode 100644 index 0000000000..1fd173ec0d --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/cron/RelationalCronJobScheduler.java @@ -0,0 +1,427 @@ +/* + * 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.scheduler.job.cron; + +import com.github.kagkarlsson.scheduler.ScheduledExecution; +import com.github.kagkarlsson.scheduler.SchedulerClient; +import com.github.kagkarlsson.scheduler.serializer.JacksonSerializer; +import com.github.kagkarlsson.scheduler.task.TaskInstance; +import com.github.kagkarlsson.scheduler.task.TaskInstanceId; +import com.google.common.base.Preconditions; +import java.time.Instant; +import java.util.ArrayList; +import java.util.Comparator; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.function.Function; +import javax.sql.DataSource; +import org.apache.gravitino.Config; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.maintenance.scheduler.AutocommitDataSources; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec.TableTarget; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpecFactory; +import org.apache.gravitino.maintenance.scheduler.TableMaintenanceSchedulerPools; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceTableRefs.ResolvedTableRef; +import org.apache.gravitino.maintenance.scheduler.policy.expand.ExpandTableCandidate; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.policy.Policy; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Schedules {@code cron-job} rows through a db-scheduler {@link SchedulerClient} on the + * entity-store DataSource. + * + * <p>Follows the usual client pattern: {@link SchedulerClient#scheduleIfNotExists} for new instance + * ids, {@link SchedulerClient#getScheduledExecutionsForTask} to inspect open rows, and {@link + * SchedulerClient#reschedule(TaskInstanceId, Instant, Object)} to replace {@code task_data} when + * filling an incomplete snapshot-expiry batch (same due time, new table list + shared jobConf). + * + * <p>The entity store disables JDBC autocommit; connections are wrapped via {@link + * AutocommitDataSources} so each client call commits. {@code task_data} is JSON via {@link + * JacksonSerializer + JsonUtils.anyFieldMapper()}. + */ +public final class RelationalCronJobScheduler implements CronJobScheduler { + + private static final Logger LOG = LoggerFactory.getLogger(RelationalCronJobScheduler.class); + + private final DataSource dataSource; + private final MaintenanceJobSpecFactory specFactory; + private final Function<List<Long>, List<ResolvedTableRef>> tableRefResolver; + + /** + * @param dataSource relational JDBC pool shared with the entity store + * @param gravitinoConfig server config for jobConf overlay + */ + public RelationalCronJobScheduler(DataSource dataSource, Config gravitinoConfig) { + this(dataSource, gravitinoConfig, MaintenanceTableRefs::resolve); + } + + /** + * Visible for tests that inject table-id resolution without a JDBC catalog store. + * + * @param dataSource relational JDBC pool + * @param gravitinoConfig server config + * @param tableRefResolver resolves entity ids to catalog + schema.table + */ + RelationalCronJobScheduler( + DataSource dataSource, + Config gravitinoConfig, + Function<List<Long>, List<ResolvedTableRef>> tableRefResolver) { + this.dataSource = Preconditions.checkNotNull(dataSource); + this.specFactory = new MaintenanceJobSpecFactory(Preconditions.checkNotNull(gravitinoConfig)); + this.tableRefResolver = Preconditions.checkNotNull(tableRefResolver); + } + + @Override + public int enqueueDueNow( + PolicyEntity policy, List<ExpandTableCandidate> candidates, int batchCapacity) { + if (policy == null || candidates == null || candidates.isEmpty()) { + return 0; + } + if (MaintenanceJobTemplates.templateName(policy.policyType()).isEmpty()) { + LOG.warn( + "Skipping cron-job enqueue for policy {} with unsupported type {}", + policy.id(), + policy.policyType()); + return 0; + } + int capacity = Math.max(1, batchCapacity); + SchedulerClient client = client(dataSource); + Instant due = Instant.now(); + + if (policy.policyType() == Policy.BuiltInType.ICEBERG_SNAPSHOT_EXPIRATION) { + return enqueueSnapshotExpiry(client, policy, candidates, capacity, due); + } + int scheduled = 0; + int limit = Math.min(candidates.size(), capacity); + for (ExpandTableCandidate candidate : candidates.subList(0, limit)) { + if (scheduleTable(client, policy, candidate.tableId(), due)) { + scheduled++; + } + } + return scheduled; + } + + /** + * Builds a SchedulerClient that can commit against the entity-store DataSource. + * + * @param dataSource entity-store pool + * @return client bound to {@link TableMaintenanceSchedulerPools#CRON_TASKS_TABLE} + */ + static SchedulerClient client(DataSource dataSource) { + return SchedulerClient.Builder.create( + AutocommitDataSources.wrap(dataSource), TableMaintenanceSchedulerPools.cronJobTask()) + .serializer(new JacksonSerializer(JsonUtils.anyFieldMapper())) + .tableName(TableMaintenanceSchedulerPools.CRON_TASKS_TABLE) + .build(); + } + + private int enqueueSnapshotExpiry( + SchedulerClient client, + PolicyEntity policy, + List<ExpandTableCandidate> slice, + int batchCapacity, + Instant due) { + List<Long> batchable = new ArrayList<>(); + int scheduled = 0; + for (ExpandTableCandidate candidate : slice) { + if (candidate.batchableForSnapshotExpiry()) { + if (!isQueuedForPolicy(client, candidate.tableId(), policy.id())) { + batchable.add(candidate.tableId()); + } + } else if (scheduleTable(client, policy, candidate.tableId(), due)) { + scheduled++; + } + } + scheduled += packIntoBatches(client, policy, batchable, batchCapacity, due); + return scheduled; + } + + private int packIntoBatches( + SchedulerClient client, + PolicyEntity policy, + List<Long> tableIds, + int batchCapacity, + Instant due) { + if (tableIds.isEmpty()) { + return 0; + } + // Same-catalog batches only: mixed catalogs would reuse wrong spark.sql.catalog.* props. + Map<String, List<Long>> byCatalog = groupByCatalog(tableIds); + int newRows = 0; + for (Map.Entry<String, List<Long>> entry : byCatalog.entrySet()) { + newRows += + packSameCatalog(client, policy, entry.getKey(), entry.getValue(), batchCapacity, due); + } + return newRows; + } + + private Map<String, List<Long>> groupByCatalog(List<Long> tableIds) { + List<ResolvedTableRef> resolved = tableRefResolver.apply(tableIds); + Map<Long, ResolvedTableRef> byId = new LinkedHashMap<>(); + for (ResolvedTableRef ref : resolved) { + byId.put(ref.tableId(), ref); + } + Map<String, List<Long>> byCatalog = new LinkedHashMap<>(); + for (Long tableId : tableIds) { + ResolvedTableRef ref = byId.get(tableId); + if (ref == null) { + LOG.warn("Skipping unresolved table {} from snapshot-expiry batch packing", tableId); + continue; + } + byCatalog.computeIfAbsent(ref.catalogName(), ignored -> new ArrayList<>()).add(tableId); + } + return byCatalog; + } + + private int packSameCatalog( + SchedulerClient client, + PolicyEntity policy, + String catalogName, + List<Long> tableIds, + int batchCapacity, + Instant due) { + List<Long> remaining = new ArrayList<>(tableIds); + int newRows = 0; + while (!remaining.isEmpty()) { + Optional<OpenBatch> incomplete = + findIncompleteBatch(client, policy.id(), batchCapacity, catalogName); + if (incomplete.isPresent()) { + OpenBatch open = incomplete.get(); + int room = batchCapacity - open.spec.tableIds().size(); + List<Long> chunk = new ArrayList<>(remaining.subList(0, Math.min(room, remaining.size()))); + if (mergeIntoBatch(client, open, chunk)) { + remaining.subList(0, chunk.size()).clear(); + continue; + } + } + int end = Math.min(batchCapacity, remaining.size()); + List<Long> chunk = new ArrayList<>(remaining.subList(0, end)); + if (scheduleBatch(client, policy, chunk, due)) { + newRows++; + remaining.subList(0, end).clear(); + continue; + } + // Fail-closed: do not drop a whole multi-table chunk on build failure — peel one id. + Long first = remaining.remove(0); + if (chunk.size() > 1 && scheduleBatch(client, policy, List.of(first), due)) { + newRows++; + } else { + LOG.warn( + "Dropping table {} from snapshot-expiry enqueue for policy {} after failed batch build", + first, + policy.id()); + } + } + return newRows; + } + + private boolean scheduleTable( + SchedulerClient client, PolicyEntity policy, long tableId, Instant due) { + Optional<MaintenanceJobSpec> spec = buildSpec(policy, List.of(tableId)); + if (spec.isEmpty()) { + return false; + } + String id = CronJobIds.forTable(tableId, policy.id()); + return client.scheduleIfNotExists( + new TaskInstance<>(TableMaintenanceSchedulerPools.CRON_JOB, id, spec.get()), due); + } + + private boolean scheduleBatch( + SchedulerClient client, PolicyEntity policy, List<Long> tableIds, Instant due) { + Optional<MaintenanceJobSpec> spec = buildSpec(policy, tableIds); + if (spec.isEmpty()) { + return false; + } + long batchId = nextBatchId(client, policy.id()); + String id = CronJobIds.forBatch(batchId, policy.id()); + return client.scheduleIfNotExists( + new TaskInstance<>(TableMaintenanceSchedulerPools.CRON_JOB, id, spec.get()), due); + } + + private boolean mergeIntoBatch(SchedulerClient client, OpenBatch open, List<Long> extraTableIds) { + Map<Long, TableTarget> byId = new LinkedHashMap<>(); + for (TableTarget target : open.spec.getTables()) { + byId.put(target.getTableId(), target); + } + String openCatalog = open.spec.getTables().get(0).catalogName(); + if (openCatalog == null || openCatalog.isEmpty()) { + LOG.warn( + "Could not merge tables into cron-job {}: open batch missing catalog name", + open.instanceId.getId()); + return false; + } + List<ResolvedTableRef> resolved = tableRefResolver.apply(extraTableIds); + // Require every requested id so packIntoBatches does not drop unresolved tables from + // remaining after a partial merge. + if (resolved.size() != extraTableIds.size()) { + LOG.warn( + "Could not merge tables into cron-job {}: resolved {} of {} ids {}", + open.instanceId.getId(), + resolved.size(), + extraTableIds.size(), + extraTableIds); + return false; + } + String metalake = open.spec.metalake(); + if (metalake == null) { + LOG.warn( + "Could not merge tables into cron-job {}: missing metalake on open spec", + open.instanceId.getId()); + return false; + } + boolean added = false; + for (ResolvedTableRef ref : resolved) { + if (!openCatalog.equals(ref.catalogName())) { + LOG.warn( + "Could not merge table {} into cron-job {}: catalog {} != open batch catalog {}", + ref.tableId(), + open.instanceId.getId(), + ref.catalogName(), + openCatalog); + return false; + } + if (!byId.containsKey(ref.tableId())) { + byId.put(ref.tableId(), MaintenanceJobSpecFactory.toTarget(metalake, ref)); + added = true; + } + } + if (!added) { + return true; + } + MaintenanceJobSpec merged = + new MaintenanceJobSpec( + new ArrayList<>(byId.values()), + open.spec.getOperations(), + open.spec.getJobConf()); + try { + return client.reschedule(open.instanceId, open.executionTime, merged); + } catch (RuntimeException e) { + LOG.debug( + "Could not merge tables into cron-job {}: {}", open.instanceId.getId(), e.toString()); + return false; + } + } + + private Optional<MaintenanceJobSpec> buildSpec(PolicyEntity policy, List<Long> tableIds) { + List<ResolvedTableRef> resolved = tableRefResolver.apply(tableIds); + if (resolved.size() != tableIds.size()) { + LOG.warn( + "Skipping cron-job enqueue for policy {}: resolved {} of {} table ids {}", + policy.id(), + resolved.size(), + tableIds.size(), + tableIds); + return Optional.empty(); + } + try { + return Optional.of(specFactory.forCron(policy, resolved)); + } catch (RuntimeException e) { + LOG.warn( + "Skipping cron-job enqueue for policy {}: failed to build jobConf / MaintenanceJobSpec: {}", + policy.id(), + e.toString()); + return Optional.empty(); + } + } + + private static Optional<OpenBatch> findIncompleteBatch( + SchedulerClient client, long policyId, int batchCapacity, String catalogName) { + return listOpenBatches(client, policyId).stream() + .filter(b -> !b.picked && b.spec.tableIds().size() < batchCapacity) + .filter(b -> catalogName.equals(b.spec.getTables().get(0).catalogName())) + .min(Comparator.comparingLong(b -> b.batchId)); + } + + private static List<OpenBatch> listOpenBatches(SchedulerClient client, long policyId) { + List<OpenBatch> open = new ArrayList<>(); + for (ScheduledExecution<MaintenanceJobSpec> execution : + client.getScheduledExecutionsForTask( + TableMaintenanceSchedulerPools.CRON_JOB, MaintenanceJobSpec.class)) { + Optional<CronJobIds.BatchPolicyIds> parsed = + CronJobIds.parseBatch(execution.getTaskInstance().getId()); + if (parsed.isEmpty() || parsed.get().policyId() != policyId) { + continue; + } + MaintenanceJobSpec data = execution.getData(); + if (data == null || !data.isValid()) { + continue; + } + open.add( + new OpenBatch( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.CRON_JOB, execution.getTaskInstance().getId()), + parsed.get().batchId(), + execution.getExecutionTime(), + data, + execution.isPicked())); + } + return open; + } + + static long nextBatchId(SchedulerClient client, long policyId) { + long max = -1L; + for (OpenBatch batch : listOpenBatches(client, policyId)) { + max = Math.max(max, batch.batchId); + } + return max + 1L; + } + + /** True when a table already has a queued {@code table:} or {@code batch:} cron-job. */ + static boolean isQueuedForPolicy(SchedulerClient client, long tableId, long policyId) { + if (client + .getScheduledExecution( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forTable(tableId, policyId))) + .isPresent()) { + return true; + } + for (OpenBatch batch : listOpenBatches(client, policyId)) { + if (batch.spec.tableIds().contains(tableId)) { + return true; + } + } + return false; + } + + private static final class OpenBatch { + private final TaskInstanceId instanceId; + private final long batchId; + private final Instant executionTime; + private final MaintenanceJobSpec spec; + private final boolean picked; + + private OpenBatch( + TaskInstanceId instanceId, + long batchId, + Instant executionTime, + MaintenanceJobSpec spec, + boolean picked) { + this.instanceId = instanceId; + this.batchId = batchId; + this.executionTime = executionTime; + this.spec = spec; + this.picked = picked; + } + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestCronJobIds.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestCronJobIds.java new file mode 100644 index 0000000000..57f89116af --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestCronJobIds.java @@ -0,0 +1,58 @@ +/* + * 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.scheduler.job.cron; + +import java.util.Optional; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class TestCronJobIds { + + @Test + void testForTableAndParseRoundTrip() { + String id = CronJobIds.forTable(10L, 20L); + Assertions.assertEquals("table:10:20", id); + Optional<CronJobIds.TablePolicyIds> parsed = CronJobIds.parseTable(id); + Assertions.assertTrue(parsed.isPresent()); + Assertions.assertEquals(10L, parsed.get().tableId()); + Assertions.assertEquals(20L, parsed.get().policyId()); + } + + @Test + void testForBatchAndParseRoundTrip() { + Assertions.assertEquals("batch:0:9", CronJobIds.forBatch(0L, 9L)); + Assertions.assertEquals("batch:12:9", CronJobIds.forBatch(12L, 9L)); + Optional<CronJobIds.BatchPolicyIds> parsed = CronJobIds.parseBatch("batch:12:9"); + Assertions.assertTrue(parsed.isPresent()); + Assertions.assertEquals(12L, parsed.get().batchId()); + Assertions.assertEquals(9L, parsed.get().policyId()); + } + + @Test + void testParseRejectsInvalid() { + Assertions.assertTrue(CronJobIds.parseTable(null).isEmpty()); + Assertions.assertTrue(CronJobIds.parseTable("batch:1:2").isEmpty()); + Assertions.assertTrue(CronJobIds.parseTable("table:").isEmpty()); + Assertions.assertTrue(CronJobIds.parseTable("table:1").isEmpty()); + Assertions.assertTrue(CronJobIds.parseTable("table:a:2").isEmpty()); + Assertions.assertTrue(CronJobIds.parseBatch(null).isEmpty()); + Assertions.assertTrue(CronJobIds.parseBatch("table:1:2").isEmpty()); + Assertions.assertTrue(CronJobIds.parseBatch("batch:pending:9").isEmpty()); + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestCronJobScheduler.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestCronJobScheduler.java new file mode 100644 index 0000000000..57a4c17596 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/cron/TestCronJobScheduler.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.scheduler.job.cron; + +import com.github.kagkarlsson.scheduler.Scheduler; +import com.github.kagkarlsson.scheduler.SchedulerClient; +import com.github.kagkarlsson.scheduler.serializer.JacksonSerializer; +import com.github.kagkarlsson.scheduler.task.TaskInstanceId; +import java.time.Duration; +import java.time.Instant; +import java.util.List; +import java.util.Map; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import javax.sql.DataSource; +import org.apache.gravitino.Config; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.maintenance.policy.MaintenanceAuthConfigs; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec; +import org.apache.gravitino.maintenance.scheduler.TableMaintenanceSchedulerPools; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceTableRefs.ResolvedTableRef; +import org.apache.gravitino.maintenance.scheduler.policy.expand.ExpandTableCandidate; +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.PolicyContents; +import org.apache.gravitino.storage.relational.TestJDBCBackend; +import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper; +import org.apache.gravitino.utils.NamespaceUtil; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.TestTemplate; +import org.apache.gravitino.maintenance.scheduler.policy.expand.ExpandOccupancyGate; + +class TestCronJobScheduler extends TestJDBCBackend { + + private static final String GRAVITINO_URI = "http://localhost:8090";; + + @TestTemplate + void testEnqueueIsIdempotentAndRespectsBatchSize() { + DataSource dataSource = dataSource(); + CronJobScheduler scheduler = scheduler(dataSource); + PolicyEntity policy = compactionPolicy(7L); + + List<ExpandTableCandidate> three = + List.of( + candidate(1L, MetadataObject.Type.TABLE), + candidate(2L, MetadataObject.Type.TABLE), + candidate(3L, MetadataObject.Type.TABLE)); + Assertions.assertEquals(2, scheduler.enqueueDueNow(policy, three, 2)); + Assertions.assertEquals( + 0, + scheduler.enqueueDueNow( + policy, + List.of( + candidate(1L, MetadataObject.Type.TABLE), candidate(2L, MetadataObject.Type.TABLE)), + 2)); + + SchedulerClient client = RelationalCronJobScheduler.client(dataSource); + Assertions.assertTrue( + client + .getScheduledExecution( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forTable(1L, 7L))) + .isPresent()); + Assertions.assertTrue( + client + .getScheduledExecution( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forTable(2L, 7L))) + .isPresent()); + Assertions.assertTrue( + client + .getScheduledExecution( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forTable(3L, 7L))) + .isEmpty()); + + MaintenanceJobSpec spec = + (MaintenanceJobSpec) + client + .getScheduledExecution( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forTable(1L, 7L))) + .get() + .getData(); + Assertions.assertTrue(spec.isValid()); + Assertions.assertTrue(spec.getJobConf().containsKey("updater_options")); + Assertions.assertEquals(List.of(1L), spec.tableIds()); + } + + @TestTemplate + void testSnapshotExpiryFillsIncompleteBatchBeforeAllocatingNextId() { + DataSource dataSource = dataSource(); + CronJobScheduler scheduler = scheduler(dataSource); + PolicyEntity policy = expirePolicy(9L); + + Assertions.assertEquals( + 1, + scheduler.enqueueDueNow( + policy, + List.of( + candidate(10L, MetadataObject.Type.SCHEMA), + candidate(11L, MetadataObject.Type.SCHEMA)), + 10)); + + Assertions.assertEquals( + 0, + scheduler.enqueueDueNow( + policy, + List.of( + candidate(10L, MetadataObject.Type.SCHEMA), + candidate(11L, MetadataObject.Type.SCHEMA), + candidate(12L, MetadataObject.Type.CATALOG), + candidate(13L, MetadataObject.Type.SCHEMA)), + 10)); + + SchedulerClient client = RelationalCronJobScheduler.client(dataSource); + TaskInstanceId batch0 = + TaskInstanceId.of(TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forBatch(0L, 9L)); + TaskInstanceId batch1 = + TaskInstanceId.of(TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forBatch(1L, 9L)); + Assertions.assertTrue(client.getScheduledExecution(batch0).isPresent()); + Assertions.assertTrue(client.getScheduledExecution(batch1).isEmpty()); + MaintenanceJobSpec data = + (MaintenanceJobSpec) client.getScheduledExecution(batch0).get().getData(); + Assertions.assertEquals(List.of(10L, 11L, 12L, 13L), data.tableIds()); + Assertions.assertEquals(1, data.getOperations().size()); + Assertions.assertTrue(data.getJobConf().containsKey("updater_options")); + } + + @TestTemplate + void testSnapshotExpiryAllocatesNextBatchWhenCapacityFull() { + DataSource dataSource = dataSource(); + CronJobScheduler scheduler = scheduler(dataSource); + PolicyEntity policy = expirePolicy(9L); + + Assertions.assertEquals( + 1, + scheduler.enqueueDueNow( + policy, + List.of( + candidate(1L, MetadataObject.Type.SCHEMA), + candidate(2L, MetadataObject.Type.SCHEMA)), + 2)); + Assertions.assertEquals( + 1, + scheduler.enqueueDueNow( + policy, + List.of( + candidate(1L, MetadataObject.Type.SCHEMA), + candidate(2L, MetadataObject.Type.SCHEMA), + candidate(3L, MetadataObject.Type.SCHEMA)), + 2)); + + SchedulerClient client = RelationalCronJobScheduler.client(dataSource); + MaintenanceJobSpec batch0 = + (MaintenanceJobSpec) + client + .getScheduledExecution( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forBatch(0L, 9L))) + .get() + .getData(); + MaintenanceJobSpec batch1 = + (MaintenanceJobSpec) + client + .getScheduledExecution( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forBatch(1L, 9L))) + .get() + .getData(); + Assertions.assertEquals(List.of(1L, 2L), batch0.tableIds()); + Assertions.assertEquals(List.of(3L), batch1.tableIds()); + } + + @TestTemplate + void testPendingCronGateSkipsQueuedTables() { + DataSource dataSource = dataSource(); + CronJobScheduler scheduler = scheduler(dataSource); + Assertions.assertEquals( + 1, + scheduler.enqueueDueNow( + expirePolicy(9L), + List.of( + candidate(10L, MetadataObject.Type.SCHEMA), + candidate(11L, MetadataObject.Type.SCHEMA)), + 10)); + + PendingCronExpandOccupancyGate gate = + new PendingCronExpandOccupancyGate((t, p, m, n) -> false, dataSource); + Assertions.assertTrue(gate.shouldSkip(10L, 9L, 0L, System.currentTimeMillis())); + Assertions.assertFalse(gate.shouldSkip(12L, 9L, 0L, System.currentTimeMillis())); + } + + @TestTemplate + void testCronJobTaskRemovesRowAfterPick() throws Exception { + DataSource dataSource = dataSource(); + CronJobScheduler enqueue = scheduler(dataSource); + Assertions.assertEquals( + 1, + enqueue.enqueueDueNow( + compactionPolicy(9L), List.of(candidate(99L, MetadataObject.Type.TABLE)), 1)); + + TaskInstanceId instanceId = + TaskInstanceId.of(TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forTable(99L, 9L)); + SchedulerClient client = RelationalCronJobScheduler.client(dataSource); + Assertions.assertTrue(client.getScheduledExecution(instanceId).isPresent()); + + Scheduler cronPool = + Scheduler.create(dataSource, TableMaintenanceSchedulerPools.cronJobTask()) + .threads(1) + .pollingInterval(Duration.ofMillis(100)) + .serializer(new JacksonSerializer(JsonUtils.anyFieldMapper())) + .tableName(TableMaintenanceSchedulerPools.CRON_TASKS_TABLE) + .commitWhenAutocommitDisabled(true) + .build(); + cronPool.start(); + try { + long deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(10); + while (client.getScheduledExecution(instanceId).isPresent()) { + Assertions.assertTrue( + System.nanoTime() < deadline, "cron-job row was not removed after pick"); + Thread.sleep(100L); + } + } finally { + cronPool.stop(); + } + + Assertions.assertEquals( + 1, + enqueue.enqueueDueNow( + compactionPolicy(9L), List.of(candidate(99L, MetadataObject.Type.TABLE)), 1)); + Assertions.assertTrue(client.getScheduledExecution(instanceId).isPresent()); + } + + private static CronJobScheduler scheduler(DataSource dataSource) { + return new RelationalCronJobScheduler( + dataSource, + configWithUri(), + ids -> + ids.stream() + .map(id -> new ResolvedTableRef(id, "iceberg", "db.t" + id)) + .collect(Collectors.toList())); + } + + private static DataSource dataSource() { + return SqlSessionFactoryHelper.getInstance() + .getSqlSessionFactory() + .getConfiguration() + .getEnvironment() + .getDataSource(); + } + + private static Config configWithUri() { + Config config = new Config(false) {}; + config.loadFromMap( + Map.of( + MaintenanceAuthConfigs.GRAVITINO_MAINTENANCE_GRAVITINO_AUTH_URI.getKey(), + GRAVITINO_URI), + t -> true); + return config; + } + + private static ExpandTableCandidate candidate(long tableId, MetadataObject.Type level) { + return new ExpandTableCandidate(tableId, level); + } + + private static PolicyEntity compactionPolicy(long id) { + return PolicyEntity.builder() + .withId(id) + .withName("p-" + id) + .withNamespace(NamespaceUtil.ofPolicy("ml")) + .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION) + .withEnabled(true) + .withContent( + 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, + "binpack", + "", + Map.of(), + null, + true, + null, + Map.of("spark_executor_memory", "4g"), + Map.of())) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.EPOCH).build()) + .build(); + } + + private static PolicyEntity expirePolicy(long id) { + return PolicyEntity.builder() + .withId(id) + .withName("e-" + id) + .withNamespace(NamespaceUtil.ofPolicy("ml")) + .withPolicyType(Policy.BuiltInType.ICEBERG_SNAPSHOT_EXPIRATION) + .withEnabled(true) + .withContent(PolicyContents.icebergSnapshotExpiration()) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.EPOCH).build()) + .build(); + } +}
