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 3217b14f70d688eab15d15a4c3d834c1d10a1736 Author: remilu <[email protected]> AuthorDate: Mon Oct 26 07:55:09 2026 +0800 [#14219] feat(tms): Enqueue commit-job specs from IRC updates Fixes #14219. Squashed from fork #80 (B10). Co-authored-by: Cursor <[email protected]> --- conf/gravitino.conf.template | 8 + design-docs/table-maintenance-service-design.md | 6 +- .../listener/PostCommitMaintenanceListener.java | 171 +++++++++++++++++ .../scheduler/TableMaintenanceSchedulerPools.java | 37 +++- .../scheduler/job/commit/CommitJobScheduler.java | 39 ++++ .../scheduler/job/commit/CommitJobSpecFactory.java | 200 ++++++++++++++++++++ .../commit/MaintenanceCommitSubmitPipeline.java | 192 ++++++++++++++++++++ .../job/commit/RelationalCommitJobScheduler.java | 202 +++++++++++++++++++++ .../mapper/TableMaintenanceActionMapper.java | 3 + .../TableMaintenanceActionSQLProviderFactory.java | 11 ++ .../TableMaintenanceActionBaseSQLProvider.java | 24 +++ .../service/TableMaintenanceActionMetaService.java | 134 ++++++++++++++ .../TestPostCommitMaintenanceListener.java | 114 ++++++++++++ .../job/commit/TestCommitJobSpecFactory.java | 172 ++++++++++++++++++ .../TestMaintenanceCommitSubmitPipeline.java | 185 +++++++++++++++++++ .../commit/TestRelationalCommitJobScheduler.java | 190 +++++++++++++++++++ .../TestTableMaintenanceActionMetaService.java | 88 ++++++++- 17 files changed, 1765 insertions(+), 11 deletions(-) diff --git a/conf/gravitino.conf.template b/conf/gravitino.conf.template index 66554eb606..f465f72f21 100644 --- a/conf/gravitino.conf.template +++ b/conf/gravitino.conf.template @@ -154,6 +154,14 @@ gravitino.lance-rest.namespace-backend = gravitino # gravitino.maintenance.snapshot-expiry.minIntervalMs = 3600000 # gravitino.maintenance.orphan-cleanup.minIntervalMs = 86400000 +# Base URI Spark jobs use for TMS sample/finish callbacks (updater_options.gravitino_uri). +# Required for cron Spark submit when the policy does not set updaterOptions.gravitino_uri. +# gravitino.maintenance.gravitinoUri = http://localhost:8090 + +# IRC → post-commit maintenance listener (§5.2.3). Requires colocated iceberg-rest aux service and TMS Feature. +# gravitino.eventListener.names = post-commit-maintenance-listener +# gravitino.eventListener.post-commit-maintenance-listener.class = org.apache.gravitino.maintenance.listener.PostCommitMaintenanceListener + # THE CONFIGURATION FOR sensitive property key matching # Comma-separated keywords. Each entry is a literal substring, not a regular expression. # This list replaces the default (secret,password,token,credential,access,account). diff --git a/design-docs/table-maintenance-service-design.md b/design-docs/table-maintenance-service-design.md index 5913c07926..da3f7e4f0b 100644 --- a/design-docs/table-maintenance-service-design.md +++ b/design-docs/table-maintenance-service-design.md @@ -473,14 +473,14 @@ where needed, mutate tables, and run Jobs. | ----------------------------------------- | ------------------------------------------------------------------------------------------ | | `gravitino.server.rest.extensionPackages` | Include TMS Feature package (e.g. `org.apache.gravitino.maintenance.web.rest.feature`) | | `gravitino.auxService.names` | Include `iceberg-rest` when using IRC (same JVM) | -| `gravitino.eventListener.names` | Include a TMS commit-job listener name (e.g. `tms-commit-job`) | +| `gravitino.eventListener.names` | Include the post-commit maintenance listener name (example below) | | `gravitino.eventListener.{name}.class` | FQCN of the TMS `EventListenerPlugin` that upserts commit-job on `IcebergUpdateTableEvent` | ```properties gravitino.server.rest.extensionPackages = org.apache.gravitino.maintenance.web.rest.feature gravitino.auxService.names = iceberg-rest -gravitino.eventListener.names = tms-commit-job -gravitino.eventListener.tms-commit-job.class = org.apache.gravitino.maintenance.listener.IcebergUpdateTableEventListener +gravitino.eventListener.names = post-commit-maintenance-listener +gravitino.eventListener.post-commit-maintenance-listener.class = org.apache.gravitino.maintenance.listener.PostCommitMaintenanceListener gravitino.maintenance.scheduler.expandThreads = 4 gravitino.maintenance.scheduler.cronThreads = 8 gravitino.maintenance.scheduler.commitThreads = 4 diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/listener/PostCommitMaintenanceListener.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/listener/PostCommitMaintenanceListener.java new file mode 100644 index 0000000000..facf3d3ec8 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/listener/PostCommitMaintenanceListener.java @@ -0,0 +1,171 @@ +/* + * 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.listener; + +import com.google.common.annotations.VisibleForTesting; +import java.util.Map; +import java.util.function.Function; +import javax.sql.DataSource; +import org.apache.gravitino.Config; +import org.apache.gravitino.GravitinoEnv; +import org.apache.gravitino.NameIdentifier; +import org.apache.gravitino.exceptions.NoSuchEntityException; +import org.apache.gravitino.listener.api.EventListenerPlugin; +import org.apache.gravitino.listener.api.event.Event; +import org.apache.gravitino.maintenance.scheduler.job.commit.CommitJobScheduler; +import org.apache.gravitino.maintenance.scheduler.job.commit.CommitJobUpsertResult; +import org.apache.gravitino.maintenance.scheduler.job.commit.RelationalCommitJobScheduler; +import org.apache.gravitino.meta.TableEntity; +import org.apache.gravitino.storage.relational.service.TableMetaService; +import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Upserts a one-shot {@code commit-job} / {@code {table_id}} after a successful IRC table update + * (§5.2.3). + * + * <p>Register with {@code gravitino.eventListener.names} / {@code .{name}.class}. Enqueue builds + * the onCommit chain spec (cooldown gate + jobConf) into {@code task_data}; the commit pool + * consumes that spec and submits Spark. + * + * <p>Uses the event class name (not a compile dependency on {@code iceberg-rest-server}) so TMS can + * load when IRC events are absent from the classpath. + */ +public class PostCommitMaintenanceListener implements EventListenerPlugin { + + private static final Logger LOG = LoggerFactory.getLogger(PostCommitMaintenanceListener.class); + + /** FQCN of the IRC post-event that wakes commit maintenance. */ + static final String ICEBERG_UPDATE_TABLE_EVENT = + "org.apache.gravitino.listener.api.event.IcebergUpdateTableEvent"; + + private CommitJobScheduler commitJobScheduler; + private Function<NameIdentifier, TableEntity> tableLookupOverride; + + @Override + public void init(Map<String, String> properties) throws RuntimeException { + // No listener-specific properties yet. + } + + @Override + public void start() throws RuntimeException { + DataSource dataSource = resolveEntityStoreDataSource(); + Config gravitinoConfig = GravitinoEnv.getInstance().config(); + if (gravitinoConfig == null) { + throw new IllegalStateException("GravitinoEnv config is not available"); + } + this.commitJobScheduler = new RelationalCommitJobScheduler(dataSource, gravitinoConfig); + LOG.info("PostCommitMaintenanceListener started (commit-job upsert with spec)"); + } + + @Override + public void stop() throws RuntimeException { + commitJobScheduler = null; + } + + @Override + public Mode mode() { + return Mode.ASYNC_ISOLATED; + } + + @Override + public void onPostEvent(Event postEvent) throws RuntimeException { + if (postEvent == null || !ICEBERG_UPDATE_TABLE_EVENT.equals(postEvent.getClass().getName())) { + return; + } + enqueueCommitJob(postEvent.identifier()); + } + + /** + * Resolves {@code identifier} to a table entity id and upserts {@code commit-job}. + * + * @param identifier Gravitino table name identifier from the IRC event + */ + @VisibleForTesting + void enqueueCommitJob(NameIdentifier identifier) { + CommitJobScheduler scheduler = commitJobScheduler; + if (scheduler == null) { + LOG.debug("Ignoring IcebergUpdateTableEvent; commit-job scheduler is not started"); + return; + } + if (identifier == null) { + return; + } + try { + long tableId = resolveTableId(identifier); + CommitJobUpsertResult result = scheduler.upsertDueNow(tableId); + switch (result) { + case INSERTED: + LOG.info("Upserted commit-job for table {} ({})", tableId, identifier); + break; + case ALREADY_QUEUED: + LOG.debug("commit-job already queued for table {} ({})", tableId, identifier); + break; + case SKIPPED: + LOG.debug( + "Skipped commit-job enqueue for table {} ({}) (empty chain, unresolved table, auth, or jobConf build)", + tableId, + identifier); + break; + default: + throw new IllegalStateException("Unexpected commit-job upsert result: " + result); + } + } catch (NoSuchEntityException e) { + LOG.debug("No Gravitino table for IRC update {}: {}", identifier, e.getMessage()); + } catch (RuntimeException e) { + LOG.warn("Failed to upsert commit-job for IRC update {}: {}", identifier, e.toString()); + } + } + + @VisibleForTesting + void setCommitJobScheduler(CommitJobScheduler commitJobScheduler) { + this.commitJobScheduler = commitJobScheduler; + } + + @VisibleForTesting + void setTableLookup(Function<NameIdentifier, TableEntity> tableLookup) { + this.tableLookupOverride = tableLookup; + } + + private long resolveTableId(NameIdentifier identifier) { + if (tableLookupOverride != null) { + TableEntity table = tableLookupOverride.apply(identifier); + if (table == null) { + throw new NoSuchEntityException("table not found: %s", identifier); + } + return table.id(); + } + TableEntity table = TableMetaService.getInstance().getTableByIdentifier(identifier); + return table.id(); + } + + private static DataSource resolveEntityStoreDataSource() { + Object dataSource = + SqlSessionFactoryHelper.getInstance() + .getSqlSessionFactory() + .getConfiguration() + .getEnvironment() + .getDataSource(); + if (!(dataSource instanceof DataSource)) { + throw new IllegalStateException("Entity store data source is not available"); + } + return (DataSource) dataSource; + } +} 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 fc1140e922..104ef7e81e 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 @@ -36,6 +36,8 @@ import org.apache.gravitino.Config; import org.apache.gravitino.json.JsonUtils; import org.apache.gravitino.maintenance.TableMaintenanceConfig; import org.apache.gravitino.maintenance.occupancy.OccupancyGate; +import org.apache.gravitino.maintenance.scheduler.job.commit.CommitOccupancyGate; +import org.apache.gravitino.maintenance.scheduler.job.commit.MaintenanceCommitSubmitPipeline; import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobConfBuilder; import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobSubmitter; import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceSparkSubmitPipeline; @@ -56,8 +58,8 @@ import org.apache.gravitino.maintenance.scheduler.policy.expand.ExpandOccupancyG * isolation; sharing one physical table across pools is unsupported (see design §5.2). * * <p>The expand pool fans a due {@code policy-expand} lease into {@code cron-job} rows with an - * enqueue-time {@link MaintenanceJobSpec}. The cron pool consumes that spec, claims a soft slot, - * and submits Spark. The commit pool still only removes its one-shot row. + * enqueue-time {@link MaintenanceJobSpec}. The cron and commit pools consume that spec, claim a + * soft slot, and submit Spark. */ public class TableMaintenanceSchedulerPools implements AutoCloseable { @@ -95,6 +97,9 @@ public class TableMaintenanceSchedulerPools implements AutoCloseable { static final AtomicReference<BiConsumer<String, MaintenanceJobSpec>> CRON_SUBMIT = new AtomicReference<>(unconfiguredSubmit("cron-job")); + /** Commit submit handler used by {@link #commitJobTask()}; bound in {@link #start}. */ + static final AtomicReference<BiConsumer<String, MaintenanceJobSpec>> COMMIT_SUBMIT = + new AtomicReference<>(unconfiguredSubmit("commit-job")); private final List<Scheduler> schedulers = new ArrayList<>(); @@ -135,6 +140,15 @@ public class TableMaintenanceSchedulerPools implements AutoCloseable { config.claimTimeoutMs()); CRON_SUBMIT.set(cronSubmit::submit); + MaintenanceCommitSubmitPipeline commitSubmit = + new MaintenanceCommitSubmitPipeline( + CommitOccupancyGate.of(occupancyGate), + TableMaintenanceActionMetaService.getInstance(), + MaintenanceJobSubmitter.fromEnv(), + new MaintenanceJobConfBuilder(gravitinoConfig), + config.claimTimeoutMs()); + COMMIT_SUBMIT.set(commitSubmit::submit); + List<Scheduler> built = List.of( buildScheduler( @@ -171,6 +185,7 @@ public class TableMaintenanceSchedulerPools implements AutoCloseable { throw new IllegalStateException("policy-expand pipeline is not configured"); }); CRON_SUBMIT.set(unconfiguredSubmit("cron-job")); + COMMIT_SUBMIT.set(unconfiguredSubmit("commit-job")); LOG.info("Stopped table maintenance scheduler pools"); } @@ -244,15 +259,27 @@ public class TableMaintenanceSchedulerPools implements AutoCloseable { } + /** - * Commit-pool task definition. Removes the one-shot row. Chained submit lands in a follow-up. + * Commit-pool task definition. Consumes the enqueue-time chain spec, claims a soft slot, submits + * Spark, then removes the one-shot row. * * @return the {@code commit-job} task */ public static Task<MaintenanceJobSpec> commitJobTask() { return Tasks.custom(COMMIT_JOB, MaintenanceJobSpec.class) .execute( - (taskInstance, executionContext) -> - (executionComplete, executionOperations) -> executionOperations.remove()); + (taskInstance, executionContext) -> { + try { + COMMIT_SUBMIT.get().accept(taskInstance.getId(), taskInstance.getData()); + } catch (RuntimeException e) { + LOG.warn( + "commit-job submit failed for instance {}; removing one-shot row", + taskInstance.getId(), + e); + } + return (executionComplete, executionOperations) -> executionOperations.remove(); + }); } + } diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/CommitJobScheduler.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/CommitJobScheduler.java new file mode 100644 index 0000000000..7d1d564c56 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/CommitJobScheduler.java @@ -0,0 +1,39 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.maintenance.scheduler.job.commit; + +/** + * Upserts one-shot {@code commit-job} rows for IRC table updates (§5.2.3). + * + * <p>Instance id is {@code {table_id}}; bursts of commits on the same table coalesce via {@code + * scheduleIfNotExists}. Enqueue builds a {@link + * org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec} when at least one onCommit type is + * due; an existing row's {@code task_data} is left untouched. + */ +public interface CommitJobScheduler { + + /** + * Builds an enqueue-time spec (chain + jobConf) and schedules a due-now {@code commit-job} when + * one is not already queued and the spec is non-empty. + * + * @param tableId table entity id + * @return insert / coalesce / skip outcome + */ + CommitJobUpsertResult upsertDueNow(long tableId); +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/CommitJobSpecFactory.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/CommitJobSpecFactory.java new file mode 100644 index 0000000000..ad83df0f4f --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/CommitJobSpecFactory.java @@ -0,0 +1,200 @@ +/* + * 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.commit; + +import com.fasterxml.jackson.core.JsonProcessingException; +import com.fasterxml.jackson.core.type.TypeReference; +import com.google.common.base.Preconditions; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.stream.Collectors; +import org.apache.gravitino.Config; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.maintenance.optimizer.common.util.IcebergSparkConfigUtils; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec.Operation; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec.TableTarget; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpecFactory; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobConfBuilder; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobTemplates; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceTableRefs.ResolvedTableRef; + +/** + * Builds enqueue-time {@link MaintenanceJobSpec} values for the commit-job path (keeps commit → + * scheduler dependency one-way; no scheduler → commit package cycle). + * + * <p>Chain {@code jobConf} starts from the first policy, takes the max of resource-shaped keys from + * later members, and {@code putIfAbsent}s remaining keys so later types keep their procedure params + * (e.g. expire {@code older_than}) when auth bags already agree. {@code rules} is unioned across + * members so each type keeps its enqueue-time rules for in-Spark recommend. + */ +final class CommitJobSpecFactory { + + private static final Set<String> CHAIN_OWNED_KEYS = + Set.of( + MaintenanceJobTemplates.MAINTENANCE_TYPES_KEY, + MaintenanceJobTemplates.POLICY_IDS_KEY, + "catalog_name", + "table_identifier", + "tables", + MaintenanceJobTemplates.RULES_KEY); + + private final MaintenanceJobConfBuilder jobConfBuilder; + + /** + * @param gravitinoConfig server config for auth overlay and gravitino URI + */ + CommitJobSpecFactory(Config gravitinoConfig) { + this.jobConfBuilder = + new MaintenanceJobConfBuilder(Preconditions.checkNotNull(gravitinoConfig)); + } + + /** + * Commit-chain spec for one table. Caller must already have filtered cooldown and verified auth + * consistency. + * + * @param chain non-empty onCommit members in order + * @param ref resolved Spark table ref + * @return spec with merged {@code jobConf} including {@code maintenance_types} and {@code + * policy_ids} + */ + MaintenanceJobSpec forCommit(List<CommitChainMember> chain, ResolvedTableRef ref) { + Preconditions.checkArgument(chain != null && !chain.isEmpty(), "chain must not be empty"); + Preconditions.checkNotNull(ref, "ref"); + + CommitChainMember first = chain.get(0); + String metalake = first.policy().namespace().level(0); + Map<String, String> jobConf = buildChainJobConf(chain, ref); + List<TableTarget> targets = List.of(MaintenanceJobSpecFactory.toTarget(metalake, ref)); + List<Operation> operations = + chain.stream() + .map(m -> new Operation(m.policy().id(), m.maintenanceType())) + .collect(Collectors.toList()); + return new MaintenanceJobSpec(targets, operations, jobConf); + } + + private Map<String, String> buildChainJobConf( + List<CommitChainMember> chain, ResolvedTableRef ref) { + CommitChainMember first = chain.get(0); + Map<String, String> jobConf = jobConfBuilder.build(first.policy(), ref.catalogName()); + for (int i = 1; i < chain.size(); i++) { + Map<String, String> other = jobConfBuilder.build(chain.get(i).policy(), ref.catalogName()); + mergeLaterMemberJobConf(jobConf, other); + } + List<String> types = + chain.stream().map(CommitChainMember::maintenanceType).collect(Collectors.toList()); + List<String> policyIds = + chain.stream().map(m -> Long.toString(m.policy().id())).collect(Collectors.toList()); + jobConf.put(MaintenanceJobTemplates.MAINTENANCE_TYPES_KEY, String.join(",", types)); + jobConf.put(MaintenanceJobTemplates.POLICY_IDS_KEY, String.join(",", policyIds)); + // Per-table keys (catalog_name / tables JSON) are applied at pick via jobConfForTable. + jobConf.remove("catalog_name"); + jobConf.remove("table_identifier"); + jobConf.remove("tables"); + return jobConf; + } + + /** + * Max-merge resource keys; fill missing non-owned keys from a later chain member (procedure + * params). First-policy auth / {@code updater_options} win. Nested {@code spark_conf} is merged + * by key with first-wins {@code putIfAbsent}. {@code rules} are unioned by maintenance type. + */ + static void mergeLaterMemberJobConf(Map<String, String> into, Map<String, String> from) { + MaintenanceJobSpecFactory.mergeResourceKeys(into, from); + mergePolicyRulesByType(into, from); + mergeSparkConfByKey(into, from); + for (Map.Entry<String, String> entry : from.entrySet()) { + String key = entry.getKey(); + String value = entry.getValue(); + if (key == null + || value == null + || CHAIN_OWNED_KEYS.contains(key) + || "spark_conf".equals(key)) { + continue; + } + into.putIfAbsent(key, value); + } + } + + /** + * Merges nested {@code spark_conf} JSON maps. Existing keys in {@code into} win; keys only in + * {@code from} are added. + */ + static void mergeSparkConfByKey(Map<String, String> into, Map<String, String> from) { + String fromJson = from.get("spark_conf"); + if (fromJson == null || fromJson.trim().isEmpty()) { + return; + } + Map<String, String> merged = + new LinkedHashMap<>( + IcebergSparkConfigUtils.parseFlatJsonMap(into.get("spark_conf"), "spark_conf")); + Map<String, String> later = IcebergSparkConfigUtils.parseFlatJsonMap(fromJson, "spark_conf"); + for (Map.Entry<String, String> entry : later.entrySet()) { + merged.putIfAbsent(entry.getKey(), entry.getValue()); + } + if (merged.isEmpty()) { + return; + } + try { + into.put("spark_conf", JsonUtils.anyFieldMapper().writeValueAsString(merged)); + } catch (JsonProcessingException e) { + throw new IllegalStateException("Failed to merge spark_conf for commit chain", e); + } + } + + /** + * Union {@link MaintenanceJobTemplates#RULES_KEY} maps so each chain type keeps its own + * enqueue-time rules for in-Spark recommend. + */ + static void mergePolicyRulesByType(Map<String, String> into, Map<String, String> from) { + String fromJson = from.get(MaintenanceJobTemplates.RULES_KEY); + if (fromJson == null || fromJson.trim().isEmpty()) { + return; + } + Map<String, Map<String, Object>> merged = + new LinkedHashMap<>(parseRulesByType(into.get(MaintenanceJobTemplates.RULES_KEY))); + Map<String, Map<String, Object>> later = parseRulesByType(fromJson); + for (Map.Entry<String, Map<String, Object>> entry : later.entrySet()) { + merged.putIfAbsent(entry.getKey(), entry.getValue()); + } + try { + into.put( + MaintenanceJobTemplates.RULES_KEY, JsonUtils.anyFieldMapper().writeValueAsString(merged)); + } catch (JsonProcessingException e) { + throw new IllegalStateException( + "Failed to serialize " + MaintenanceJobTemplates.RULES_KEY, e); + } + } + + private static Map<String, Map<String, Object>> parseRulesByType(String json) { + if (json == null || json.trim().isEmpty()) { + return Map.of(); + } + try { + Map<String, Map<String, Object>> parsed = + JsonUtils.anyFieldMapper() + .readValue(json, new TypeReference<Map<String, Map<String, Object>>>() {}); + return parsed == null ? Map.of() : parsed; + } catch (JsonProcessingException e) { + throw new IllegalStateException("Failed to parse " + MaintenanceJobTemplates.RULES_KEY, e); + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/MaintenanceCommitSubmitPipeline.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/MaintenanceCommitSubmitPipeline.java new file mode 100644 index 0000000000..91a7d5df1d --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/MaintenanceCommitSubmitPipeline.java @@ -0,0 +1,192 @@ +/* + * 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.commit; + +import com.google.common.base.Preconditions; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.stream.Collectors; +import javax.annotation.Nullable; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec.Operation; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec.TableTarget; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpecFactory; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobConfBuilder; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobSubmitter; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobTemplates; +import org.apache.gravitino.maintenance.storage.service.OccupancyClaim; +import org.apache.gravitino.maintenance.storage.service.TableMaintenanceActionMetaService; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Commit-pool control-plane submit (§5.2.3): consume enqueue-time {@link MaintenanceJobSpec} → + * abort if any step is in-flight → soft-slot claim for every chain policy → one chained {@code + * runJob} → bind the same job id on every claim. + * + * <p>Chain collection, cooldown filtering, auth consistency, and jobConf merge happen at enqueue + * (see {@link RelationalCommitJobScheduler}). The caller always deletes the one-shot {@code + * commit-job} row after this method returns. + */ +public final class MaintenanceCommitSubmitPipeline { + + private static final Logger LOG = LoggerFactory.getLogger(MaintenanceCommitSubmitPipeline.class); + + private final CommitOccupancyGate occupancyGate; + private final TableMaintenanceActionMetaService actionMetaService; + private final MaintenanceJobSubmitter jobSubmitter; + private final MaintenanceJobConfBuilder jobConfBuilder; + private final long claimTimeoutMs; + + /** + * @param occupancyGate in-flight checks + * @param actionMetaService soft-slot claim / bind + * @param jobSubmitter Gravitino Job submitter + * @param jobConfBuilder resolves enqueue-time conf-key refs before {@code runJob} + * @param claimTimeoutMs unbound claim TTL + */ + public MaintenanceCommitSubmitPipeline( + CommitOccupancyGate occupancyGate, + TableMaintenanceActionMetaService actionMetaService, + MaintenanceJobSubmitter jobSubmitter, + MaintenanceJobConfBuilder jobConfBuilder, + long claimTimeoutMs) { + this.occupancyGate = Preconditions.checkNotNull(occupancyGate); + this.actionMetaService = Preconditions.checkNotNull(actionMetaService); + this.jobSubmitter = Preconditions.checkNotNull(jobSubmitter); + this.jobConfBuilder = Preconditions.checkNotNull(jobConfBuilder); + Preconditions.checkArgument(claimTimeoutMs > 0, "claimTimeoutMs must be positive"); + this.claimTimeoutMs = claimTimeoutMs; + } + + /** + * Runs one commit-job pick for {@code instanceId} ({@code {table_id}}). + * + * @param instanceId db-scheduler commit-job instance id + * @param spec enqueue-time payload; required + */ + public void submit(String instanceId, @Nullable MaintenanceJobSpec spec) { + long tableId; + try { + tableId = Long.parseLong(instanceId); + } catch (NumberFormatException e) { + LOG.warn("Skipping commit-job submit for unparseable instance id {}", instanceId); + return; + } + if (tableId <= 0) { + LOG.warn("Skipping commit-job submit for non-positive table id {}", instanceId); + return; + } + if (spec == null || !spec.isValid()) { + LOG.warn( + "Skipping commit-job submit for table {} with missing/invalid task_data spec", tableId); + return; + } + TableTarget target = spec.getTables().get(0); + if (target.getTableId() != tableId) { + LOG.warn( + "Skipping commit-job submit for instance {} whose spec tableId is {}", + instanceId, + target.getTableId()); + return; + } + + // Design §5.2.3: any remaining type in-flight → delete commit-job and stop. + for (Operation step : spec.getOperations()) { + if (occupancyGate.isInFlight(tableId, step.getPolicyId())) { + LOG.info( + "Skipping commit-job submit for table {} (type {} policy {} still in-flight)", + tableId, + step.getMaintenanceType(), + step.getPolicyId()); + return; + } + } + + List<Long> policyIds = + spec.getOperations().stream().map(Operation::getPolicyId).collect(Collectors.toList()); + Optional<List<OccupancyClaim>> claims = + actionMetaService.tryAcquireInFlightBatch(tableId, policyIds, claimTimeoutMs); + if (claims.isEmpty()) { + LOG.info( + "Skipping commit-job submit for table {} (lost soft-slot race on chain policies {})", + tableId, + policyIds); + return; + } + + runAndBind(spec, claims.get(), tableId); + } + + private void runAndBind(MaintenanceJobSpec spec, List<OccupancyClaim> claims, long tableId) { + String metalake = spec.metalake(); + Map<String, String> jobConf; + try { + jobConf = + jobConfBuilder.resolveForSubmit( + MaintenanceJobSpecFactory.jobConfForTable( + spec.getJobConf(), spec.getTables().get(0))); + } catch (RuntimeException e) { + actionMetaService.deleteClaims(claims); + LOG.warn( + "Skipping commit-job submit for table {}: invalid spec table target: {}", + tableId, + e.toString()); + return; + } + Long jobId = null; + try { + jobId = + jobSubmitter.runJob( + metalake, MaintenanceJobTemplates.TMS_ICEBERG_MAINTENANCE_TEMPLATE, jobConf); + actionMetaService.bindJobIds(claims, jobId); + LOG.info( + "Submitted commit-job for table {} as job {} (chain={}, policies={})", + tableId, + jobId, + spec.getOperations().stream() + .map(Operation::getMaintenanceType) + .collect(Collectors.toList()), + spec.getOperations().stream().map(Operation::getPolicyId).collect(Collectors.toList())); + } catch (RuntimeException e) { + if (jobId != null) { + cancelJobQuietly(metalake, jobId); + } + actionMetaService.deleteClaims(claims); + LOG.warn( + "Failed commit-job submit for table {} policies={}: {}", + tableId, + spec.getOperations().stream().map(Operation::getPolicyId).collect(Collectors.toList()), + e.toString()); + } + } + + private void cancelJobQuietly(String metalake, long jobId) { + try { + jobSubmitter.cancelJob(metalake, jobId); + } catch (RuntimeException cancelError) { + LOG.warn( + "Failed to cancel orphaned commit Spark job {} under metalake {} after bind failure", + jobId, + metalake, + cancelError); + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/RelationalCommitJobScheduler.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/RelationalCommitJobScheduler.java new file mode 100644 index 0000000000..4eb5524f88 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/job/commit/RelationalCommitJobScheduler.java @@ -0,0 +1,202 @@ +/* + * 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.commit; + +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.List; +import java.util.Optional; +import java.util.function.Function; +import javax.annotation.Nullable; +import javax.sql.DataSource; +import org.apache.gravitino.Config; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.maintenance.TableMaintenanceConfig; +import org.apache.gravitino.maintenance.occupancy.OccupancyGate; +import org.apache.gravitino.maintenance.policy.MaintenanceMinIntervalConfigs; +import org.apache.gravitino.maintenance.scheduler.AutocommitDataSources; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec; +import org.apache.gravitino.maintenance.scheduler.TableMaintenanceSchedulerPools; +import org.apache.gravitino.maintenance.storage.service.TableMaintenanceActionMetaService; +import org.apache.gravitino.policy.ExecutablePolicyContent; +import org.apache.gravitino.policy.PolicyContent; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Schedules {@code commit-job} rows through a db-scheduler {@link SchedulerClient} on the + * entity-store DataSource. + * + * <p>Enqueue collects the onCommit chain, drops members still inside {@code minIntervalMs}, checks + * auth consistency, builds a {@link MaintenanceJobSpec}, and inserts only when the spec is + * non-empty. Coalesce leaves an existing row's {@code task_data} untouched. + */ +public final class RelationalCommitJobScheduler implements CommitJobScheduler { + + private static final Logger LOG = LoggerFactory.getLogger(RelationalCommitJobScheduler.class); + + private final DataSource dataSource; + private final Config gravitinoConfig; + private final Function<Long, Optional<CollectedCommitChain>> chainCollector; + private final CommitOccupancyGate occupancyGate; + private final CommitJobSpecFactory specFactory; + + /** + * @param dataSource relational JDBC pool shared with the entity store + * @param gravitinoConfig server config for auth overlay and minInterval fallbacks + */ + public RelationalCommitJobScheduler(DataSource dataSource, Config gravitinoConfig) { + this( + dataSource, + gravitinoConfig, + new CommitChainCollector()::collect, + CommitOccupancyGate.of( + new OccupancyGate( + TableMaintenanceActionMetaService.getInstance(), + TableMaintenanceConfig.from(gravitinoConfig).claimTimeoutMs()))); + } + + /** + * Visible for tests. + * + * @param dataSource JDBC pool + * @param gravitinoConfig server config + * @param chainCollector returns the collected onCommit chain for a table id + * @param occupancyGate cooldown / in-flight checks for enqueue gating + */ + RelationalCommitJobScheduler( + DataSource dataSource, + Config gravitinoConfig, + Function<Long, Optional<CollectedCommitChain>> chainCollector, + CommitOccupancyGate occupancyGate) { + this.dataSource = Preconditions.checkNotNull(dataSource); + this.gravitinoConfig = Preconditions.checkNotNull(gravitinoConfig); + this.chainCollector = Preconditions.checkNotNull(chainCollector); + this.occupancyGate = Preconditions.checkNotNull(occupancyGate); + this.specFactory = new CommitJobSpecFactory(gravitinoConfig); + } + + @Override + public CommitJobUpsertResult upsertDueNow(long tableId) { + Preconditions.checkArgument(tableId > 0, "tableId must be positive"); + SchedulerClient client = client(dataSource); + String id = Long.toString(tableId); + // Coalesce without rebuilding task_data: keep the first enqueue-time spec. + if (client + .getScheduledExecution(TaskInstanceId.of(TableMaintenanceSchedulerPools.COMMIT_JOB, id)) + .isPresent()) { + return CommitJobUpsertResult.ALREADY_QUEUED; + } + Optional<MaintenanceJobSpec> spec = buildSpec(tableId); + if (spec.isEmpty()) { + return CommitJobUpsertResult.SKIPPED; + } + boolean inserted = + client.scheduleIfNotExists( + new TaskInstance<>(TableMaintenanceSchedulerPools.COMMIT_JOB, id, spec.get()), + Instant.now()); + return inserted ? CommitJobUpsertResult.INSERTED : CommitJobUpsertResult.ALREADY_QUEUED; + } + + /** + * Builds a SchedulerClient that can commit against the entity-store DataSource. + * + * @param dataSource entity-store pool + * @return client bound to {@link TableMaintenanceSchedulerPools#COMMIT_TASKS_TABLE} + */ + static SchedulerClient client(DataSource dataSource) { + return SchedulerClient.Builder.create( + AutocommitDataSources.wrap(dataSource), TableMaintenanceSchedulerPools.commitJobTask()) + .serializer(new JacksonSerializer(JsonUtils.anyFieldMapper())) + .tableName(TableMaintenanceSchedulerPools.COMMIT_TASKS_TABLE) + .build(); + } + + private Optional<MaintenanceJobSpec> buildSpec(long tableId) { + Optional<CollectedCommitChain> collected = chainCollector.apply(tableId); + if (collected.isEmpty()) { + return Optional.empty(); + } + List<CommitChainMember> members = collected.get().members(); + if (members.isEmpty()) { + LOG.debug("Skipping commit-job enqueue for table {}: empty onCommit chain", tableId); + return Optional.empty(); + } + + long nowMs = System.currentTimeMillis(); + List<CommitChainMember> chain = filterDueToMinInterval(tableId, members, nowMs); + if (chain.isEmpty()) { + LOG.debug( + "Skipping commit-job enqueue for table {}: all onCommit types within minIntervalMs", + tableId); + return Optional.empty(); + } + + try { + CommitAuthConsistency.requireConsistent(chain, gravitinoConfig); + } catch (IllegalArgumentException e) { + LOG.error("Skipping commit-job enqueue for table {}: {}", tableId, e.getMessage()); + return Optional.empty(); + } + + try { + return Optional.of(specFactory.forCommit(chain, collected.get().tableRef())); + } catch (RuntimeException e) { + LOG.warn( + "Skipping commit-job enqueue for table {}: failed to build jobConf spec: {}", + tableId, + e.toString()); + return Optional.empty(); + } + } + + private List<CommitChainMember> filterDueToMinInterval( + long tableId, List<CommitChainMember> collected, long nowMs) { + List<CommitChainMember> due = new ArrayList<>(); + for (CommitChainMember member : collected) { + long minIntervalMs = + MaintenanceMinIntervalConfigs.resolve( + gravitinoConfig, member.policyType(), policyMinIntervalMs(member.policy().content())); + if (occupancyGate.withinMinInterval(tableId, member.policy().id(), minIntervalMs, nowMs)) { + LOG.debug( + "Dropping type {} policy {} from commit-job enqueue for table {} (within minIntervalMs={})", + member.policyType(), + member.policy().id(), + tableId, + minIntervalMs); + continue; + } + due.add(member); + } + return due; + } + + @Nullable + private static Long policyMinIntervalMs(@Nullable PolicyContent content) { + if (!(content instanceof ExecutablePolicyContent)) { + return null; + } + return ((ExecutablePolicyContent) content).executionSpec().minIntervalMs(); + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionMapper.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionMapper.java index 0453d0cd0a..522c4dbcc2 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionMapper.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionMapper.java @@ -37,6 +37,9 @@ public interface TableMaintenanceActionMapper { @Options(useGeneratedKeys = true, keyProperty = "actionMeta.id", keyColumn = "id") void insertAction(@Param("actionMeta") TableMaintenanceActionPO actionMeta); + @InsertProvider(type = TableMaintenanceActionSQLProviderFactory.class, method = "insertActions") + void insertActions(@Param("tableId") long tableId, @Param("policyIds") List<Long> policyIds); + @UpdateProvider( type = TableMaintenanceActionSQLProviderFactory.class, method = "reclaimExpiredUnbound") diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionSQLProviderFactory.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionSQLProviderFactory.java index f4a0aaf3fe..8a9818b84e 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionSQLProviderFactory.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/TableMaintenanceActionSQLProviderFactory.java @@ -19,6 +19,7 @@ package org.apache.gravitino.maintenance.storage.mapper; import com.google.common.collect.ImmutableMap; +import java.util.List; import java.util.Map; import org.apache.gravitino.maintenance.storage.mapper.provider.base.TableMaintenanceActionBaseSQLProvider; import org.apache.gravitino.maintenance.storage.mapper.provider.h2.TableMaintenanceActionH2Provider; @@ -56,6 +57,16 @@ public class TableMaintenanceActionSQLProviderFactory { return getProvider().insertAction(actionMeta); } + /** + * @param tableId table entity id + * @param policyIds policy entity ids + * @return multi-row INSERT statement + */ + public static String insertActions( + @Param("tableId") long tableId, @Param("policyIds") List<Long> policyIds) { + return getProvider().insertActions(tableId, policyIds); + } + /** * @param tableId table entity id * @param policyId policy entity id diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/provider/base/TableMaintenanceActionBaseSQLProvider.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/provider/base/TableMaintenanceActionBaseSQLProvider.java index e29ca6c454..5da601c0e7 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/provider/base/TableMaintenanceActionBaseSQLProvider.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/mapper/provider/base/TableMaintenanceActionBaseSQLProvider.java @@ -20,6 +20,7 @@ package org.apache.gravitino.maintenance.storage.mapper.provider.base; import static org.apache.gravitino.maintenance.storage.mapper.TableMaintenanceActionMapper.TABLE_NAME; +import java.util.List; import org.apache.gravitino.maintenance.storage.po.TableMaintenanceActionPO; import org.apache.gravitino.storage.relational.mapper.provider.DatabaseTimeSQL; import org.apache.ibatis.annotations.Param; @@ -51,6 +52,29 @@ public class TableMaintenanceActionBaseSQLProvider { + ", 0)"; } + /** + * Inserts unbound occupancy claims for a commit chain in one statement. Each {@code (table_id, + * policy_id, finished_at=0)} must be free or the unique index rejects the whole batch. + * + * @param tableId table entity id (shared by all chain members) + * @param policyIds policy entity ids in chain order + * @return multi-row INSERT statement + */ + public String insertActions( + @Param("tableId") long tableId, @Param("policyIds") List<Long> policyIds) { + String now = currentTimeMillisExpr(); + return "<script>" + + "INSERT INTO " + + TABLE_NAME + + " (table_id, policy_id, created_at, finished_at) VALUES " + + "<foreach item='policyId' collection='policyIds' separator=','>" + + "(#{tableId}, #{policyId}, " + + now + + ", 0)" + + "</foreach>" + + "</script>"; + } + /** * Extends an expired unbound claim lease by CAS on {@code created_at} (ShedLock-style). Lease age * is evaluated against the database clock so every node shares one TTL timeline. diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableMaintenanceActionMetaService.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableMaintenanceActionMetaService.java index f73ebfd59c..4451b57cf9 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableMaintenanceActionMetaService.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/storage/service/TableMaintenanceActionMetaService.java @@ -22,7 +22,9 @@ import static org.apache.gravitino.metrics.source.MetricsSource.GRAVITINO_RELATI import com.google.common.base.Preconditions; import java.sql.SQLIntegrityConstraintViolationException; +import java.util.ArrayList; import java.util.Collections; +import java.util.LinkedHashSet; import java.util.List; import java.util.Locale; import java.util.Objects; @@ -120,6 +122,78 @@ public class TableMaintenanceActionMetaService { return Optional.of(new OccupancyClaim(inserted.id(), inserted.createdAt())); } + /** + * Tries to acquire unbound soft slots for every entry in {@code policyIds} in one transaction + * (commit-chain submit). + * + * <p>Expired unbound rows are reclaimed first; remaining policies are inserted with a single + * multi-row {@code INSERT}. Any live slot or unique conflict aborts the transaction so the caller + * never holds a partial chain. + * + * @param tableId table entity id + * @param policyIds chain member policy ids (duplicates collapsed, order preserved) + * @param claimTimeoutMs unbound claim TTL + * @return one {@link OccupancyClaim} per distinct policy id, or empty when the chain cannot be + * claimed atomically + */ + @Monitored( + metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME, + baseMetricName = "tryAcquireTableMaintenanceActionBatch") + public Optional<List<OccupancyClaim>> tryAcquireInFlightBatch( + long tableId, List<Long> policyIds, long claimTimeoutMs) { + Preconditions.checkArgument(claimTimeoutMs > 0, "claimTimeoutMs must be positive"); + Preconditions.checkArgument( + policyIds != null && !policyIds.isEmpty(), "policyIds must not be empty"); + List<Long> distinct = new ArrayList<>(new LinkedHashSet<>(policyIds)); + try { + return Optional.of( + SessionUtils.doWithCommitAndFetchResult( + TableMaintenanceActionMapper.class, + mapper -> acquireBatchInTransaction(mapper, tableId, distinct, claimTimeoutMs))); + } catch (SoftSlotBusyException e) { + return Optional.empty(); + } catch (RuntimeException e) { + if (isUniqueViolation(e)) { + return Optional.empty(); + } + throw e; + } + } + + private static List<OccupancyClaim> acquireBatchInTransaction( + TableMaintenanceActionMapper mapper, + long tableId, + List<Long> policyIds, + long claimTimeoutMs) { + List<Long> toInsert = new ArrayList<>(); + for (long policyId : policyIds) { + int reclaimed = mapper.reclaimExpiredUnbound(tableId, policyId, claimTimeoutMs); + if (reclaimed == 1) { + continue; + } + TableMaintenanceActionPO live = + mapper.selectLiveUnboundByTableAndPolicy(tableId, policyId, claimTimeoutMs); + if (live != null) { + throw new SoftSlotBusyException(); + } + toInsert.add(policyId); + } + if (!toInsert.isEmpty()) { + mapper.insertActions(tableId, toInsert); + } + List<OccupancyClaim> claims = new ArrayList<>(policyIds.size()); + for (long policyId : policyIds) { + TableMaintenanceActionPO unbound = mapper.selectUnboundByTableAndPolicy(tableId, policyId); + Preconditions.checkState( + unbound != null && unbound.id() != null && unbound.createdAt() != null, + "batch acquire lost unbound row for table %s policy %s", + tableId, + policyId); + claims.add(new OccupancyClaim(unbound.id(), unbound.createdAt())); + } + return claims; + } + /** * Acquires an unbound claim using {@link TableMaintenanceConfig#claimTimeoutMs()} or fails. * @@ -227,6 +301,61 @@ public class TableMaintenanceActionMetaService { } } + /** + * Binds the same Gravitino job run id to every claim from {@link #tryAcquireInFlightBatch} in one + * transaction (commit-chain submit). + * + * <p>Idempotent per claim when the row is already bound to {@code jobId} and still in-flight + * (Spark may lazy-bind before submit-side bind), matching {@link #bindJobId}. + * + * @param claims claims returned by batch acquire + * @param jobId Gravitino job run id + */ + @Monitored( + metricsSource = GRAVITINO_RELATIONAL_STORE_METRIC_NAME, + baseMetricName = "bindTableMaintenanceActionJobIds") + public void bindJobIds(List<OccupancyClaim> claims, long jobId) { + Preconditions.checkArgument(claims != null && !claims.isEmpty(), "claims must not be empty"); + SessionUtils.doWithCommit( + TableMaintenanceActionMapper.class, + mapper -> { + for (OccupancyClaim claim : claims) { + Preconditions.checkNotNull(claim, "claim must not be null"); + int updated = mapper.bindJobId(claim.id(), jobId, claim.createdAt()); + if (updated == 1) { + continue; + } + TableMaintenanceActionPO current = mapper.selectById(claim.id()); + Preconditions.checkState( + current != null + && Objects.equals(current.jobId(), jobId) + && current.finishedAt() != null + && current.finishedAt() == 0L, + "expected to bind job_id on claim id=%s, updated=%s", + claim.id(), + updated); + } + }); + } + + /** + * Releases unbound claims after a failed submit. Each delete is fenced on {@code (id, + * created_at)} so a concurrent reclaim's live claim is not removed. + * + * @param claims claims to release + */ + public void deleteClaims(List<OccupancyClaim> claims) { + if (claims == null || claims.isEmpty()) { + return; + } + for (OccupancyClaim claim : claims) { + if (claim == null) { + continue; + } + deleteUnboundClaim(claim); + } + } + private static OccupancyClaim readUnboundClaim(long tableId, long policyId) { TableMaintenanceActionPO unbound = SessionUtils.getWithoutCommit( @@ -534,4 +663,9 @@ public class TableMaintenanceActionMetaService { } return false; } + + /** Thrown inside a batch-acquire transaction to roll back when any soft slot is busy. */ + private static final class SoftSlotBusyException extends RuntimeException { + private static final long serialVersionUID = 1L; + } } diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/listener/TestPostCommitMaintenanceListener.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/listener/TestPostCommitMaintenanceListener.java new file mode 100644 index 0000000000..e635ce66f7 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/listener/TestPostCommitMaintenanceListener.java @@ -0,0 +1,114 @@ +/* + * 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.listener; + +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.time.Instant; +import org.apache.gravitino.NameIdentifier; +import org.apache.gravitino.listener.api.event.Event; +import org.apache.gravitino.listener.api.event.OperationStatus; +import org.apache.gravitino.listener.api.event.OperationType; +import org.apache.gravitino.maintenance.scheduler.job.commit.CommitJobScheduler; +import org.apache.gravitino.maintenance.scheduler.job.commit.CommitJobUpsertResult; +import org.apache.gravitino.meta.AuditInfo; +import org.apache.gravitino.meta.TableEntity; +import org.apache.gravitino.utils.NamespaceUtil; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class TestPostCommitMaintenanceListener { + + @Test + void testExpectedEventClassNameConstant() { + Assertions.assertEquals( + "org.apache.gravitino.listener.api.event.IcebergUpdateTableEvent", + PostCommitMaintenanceListener.ICEBERG_UPDATE_TABLE_EVENT); + } + + @Test + void testIgnoresUnrelatedEvents() { + PostCommitMaintenanceListener listener = new PostCommitMaintenanceListener(); + CommitJobScheduler scheduler = mock(CommitJobScheduler.class); + listener.setCommitJobScheduler(scheduler); + + listener.onPostEvent(new OtherEvent("u", NameIdentifier.of("ml", "c", "s", "t"))); + + verify(scheduler, never()).upsertDueNow(anyLong()); + } + + @Test + void testEnqueueUpsertsCommitJob() { + PostCommitMaintenanceListener listener = new PostCommitMaintenanceListener(); + CommitJobScheduler scheduler = mock(CommitJobScheduler.class); + when(scheduler.upsertDueNow(eq(42L))).thenReturn(CommitJobUpsertResult.INSERTED); + listener.setCommitJobScheduler(scheduler); + listener.setTableLookup( + ident -> + TableEntity.builder() + .withId(42L) + .withName("t") + .withNamespace(NamespaceUtil.ofTable("ml", "c", "s")) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.EPOCH).build()) + .build()); + + listener.enqueueCommitJob(NameIdentifier.of("ml", "c", "s", "t")); + + verify(scheduler).upsertDueNow(42L); + } + + @Test + void testEnqueueSkippedWhenSchedulerNotStarted() { + PostCommitMaintenanceListener listener = new PostCommitMaintenanceListener(); + listener.setTableLookup( + ident -> + TableEntity.builder() + .withId(1L) + .withName("t") + .withNamespace(NamespaceUtil.ofTable("ml", "c", "s")) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.EPOCH).build()) + .build()); + + Assertions.assertDoesNotThrow( + () -> listener.enqueueCommitJob(NameIdentifier.of("ml", "c", "s", "t"))); + } + + private static final class OtherEvent extends Event { + OtherEvent(String user, NameIdentifier identifier) { + super(user, identifier); + } + + @Override + public OperationType operationType() { + return OperationType.CREATE_TABLE; + } + + @Override + public OperationStatus operationStatus() { + return OperationStatus.SUCCESS; + } + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/commit/TestCommitJobSpecFactory.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/commit/TestCommitJobSpecFactory.java new file mode 100644 index 0000000000..c6bb36e038 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/commit/TestCommitJobSpecFactory.java @@ -0,0 +1,172 @@ +/* + * 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.commit; + +import java.time.Instant; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import org.apache.gravitino.Config; +import org.apache.gravitino.json.JsonUtils; +import org.apache.gravitino.maintenance.optimizer.common.util.IcebergSparkConfigUtils; +import org.apache.gravitino.maintenance.policy.MaintenanceAuthConfigs; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobTemplates; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceTableRefs.ResolvedTableRef; +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.utils.NamespaceUtil; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class TestCommitJobSpecFactory { + + private static final String GRAVITINO_URI = "http://localhost:8090";; + + @Test + void testForCommitMergesChainAndSetsMaintenanceTypes() { + CommitJobSpecFactory factory = new CommitJobSpecFactory(configWithUri()); + CommitChainMember compaction = + new CommitChainMember( + Policy.BuiltInType.ICEBERG_COMPACTION, compaction(202L), "rewrite_data_files"); + CommitChainMember expire = + new CommitChainMember( + Policy.BuiltInType.ICEBERG_SNAPSHOT_EXPIRATION, expire(303L), "expire_snapshots"); + MaintenanceJobSpec spec = + factory.forCommit( + List.of(compaction, expire), new ResolvedTableRef(101L, "iceberg", "db.t")); + + Assertions.assertTrue(spec.isValid()); + Assertions.assertEquals(2, spec.getOperations().size()); + Assertions.assertEquals( + "rewrite_data_files,expire_snapshots", + spec.getJobConf().get(MaintenanceJobTemplates.MAINTENANCE_TYPES_KEY)); + Assertions.assertEquals( + "202,303", spec.getJobConf().get(MaintenanceJobTemplates.POLICY_IDS_KEY)); + Assertions.assertFalse(spec.getJobConf().containsKey("tms_table_id")); + Assertions.assertFalse(spec.getJobConf().containsKey("tms_policy_id")); + Assertions.assertNull(spec.getJobConf().get("catalog_name")); + Assertions.assertNull(spec.getJobConf().get("table_identifier")); + Assertions.assertNull(spec.getJobConf().get("tables")); + Assertions.assertEquals(101L, spec.getTables().get(0).getTableId()); + Assertions.assertEquals(202L, spec.getOperations().get(0).getPolicyId()); + // Later chain members contribute procedure params the first policy does not set. + Assertions.assertEquals("7d", spec.getJobConf().get("older_than")); + Assertions.assertEquals("10", spec.getJobConf().get("retain_last")); + Assertions.assertEquals("4g", spec.getJobConf().get("spark_executor_memory")); + String rulesByType = spec.getJobConf().get(MaintenanceJobTemplates.RULES_KEY); + Assertions.assertNotNull(rulesByType); + Assertions.assertTrue(rulesByType.contains("\"rewrite_data_files\"")); + Assertions.assertTrue(rulesByType.contains("\"expire_snapshots\"")); + Assertions.assertTrue(rulesByType.contains("trigger-expr")); + } + + @Test + void testMergeLaterMemberJobConfMergesSparkConfByKey() { + Map<String, String> into = new HashMap<>(); + into.put("spark_conf", "{\"spark.executor.memory\":\"4g\",\"shared\":\"first\"}"); + Map<String, String> from = + Map.of( + "spark_conf", + "{\"spark.executor.cores\":\"2\",\"shared\":\"later\",\"spark.driver.memory\":\"2g\"}"); + + CommitJobSpecFactory.mergeLaterMemberJobConf(into, from); + + Map<String, String> merged = + IcebergSparkConfigUtils.parseFlatJsonMap(into.get("spark_conf"), "spark_conf"); + Assertions.assertEquals("4g", merged.get("spark.executor.memory")); + Assertions.assertEquals("first", merged.get("shared")); + Assertions.assertEquals("2", merged.get("spark.executor.cores")); + Assertions.assertEquals("2g", merged.get("spark.driver.memory")); + } + + @Test + void testMergeLaterMemberJobConfAddsSparkConfWhenMissing() throws Exception { + Map<String, String> into = new HashMap<>(); + into.put("older_than", "7d"); + Map<String, String> from = + Map.of("spark_conf", JsonUtils.anyFieldMapper().writeValueAsString(Map.of("k", "v"))); + + CommitJobSpecFactory.mergeLaterMemberJobConf(into, from); + + Assertions.assertEquals("7d", into.get("older_than")); + Assertions.assertEquals( + "v", + IcebergSparkConfigUtils.parseFlatJsonMap(into.get("spark_conf"), "spark_conf").get("k")); + } + + 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 PolicyEntity compaction(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 expire(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( + null, + true, + null, + Map.of("older_than", "7d", "retain_last", "10", "spark_executor_memory", "2g"), + Map.of())) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.EPOCH).build()) + .build(); + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/commit/TestMaintenanceCommitSubmitPipeline.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/commit/TestMaintenanceCommitSubmitPipeline.java new file mode 100644 index 0000000000..296e129d7c --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/commit/TestMaintenanceCommitSubmitPipeline.java @@ -0,0 +1,185 @@ +/* + * 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.commit; + +import static org.mockito.ArgumentMatchers.anyList; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.anyMap; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +import java.util.HashSet; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import org.apache.gravitino.Config; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec.Operation; +import org.apache.gravitino.maintenance.scheduler.MaintenanceJobSpec.TableTarget; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobConfBuilder; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobSubmitter; +import org.apache.gravitino.maintenance.scheduler.job.cron.MaintenanceJobTemplates; +import org.apache.gravitino.maintenance.storage.service.OccupancyClaim; +import org.apache.gravitino.maintenance.storage.service.TableMaintenanceActionMetaService; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; +import org.mockito.ArgumentCaptor; + +class TestMaintenanceCommitSubmitPipeline { + + private static final long TABLE_ID = 101L; + private static final long POLICY_ID = 202L; + private static final long EXPIRE_POLICY_ID = 303L; + private static final long CLAIM_TIMEOUT_MS = 60_000L; + + @Test + void testUnparseableInstanceSkipsSubmit() { + MaintenanceJobSubmitter submitter = mock(MaintenanceJobSubmitter.class); + TableMaintenanceActionMetaService actions = mock(TableMaintenanceActionMetaService.class); + pipeline(allowAllGate(), actions, submitter).submit("not-a-table-id", spec()); + verify(submitter, never()).runJob(anyString(), anyString(), anyMap()); + } + + @Test + void testNullSpecSkipsSubmit() { + MaintenanceJobSubmitter submitter = mock(MaintenanceJobSubmitter.class); + TableMaintenanceActionMetaService actions = mock(TableMaintenanceActionMetaService.class); + pipeline(allowAllGate(), actions, submitter).submit(Long.toString(TABLE_ID), null); + verify(submitter, never()).runJob(anyString(), anyString(), anyMap()); + verify(actions, never()).tryAcquireInFlightBatch(anyLong(), anyList(), anyLong()); + } + + @Test + void testInFlightRemainingMemberSkipsSubmit() { + MaintenanceJobSubmitter submitter = mock(MaintenanceJobSubmitter.class); + TableMaintenanceActionMetaService actions = mock(TableMaintenanceActionMetaService.class); + pipeline(inFlightGate(POLICY_ID), actions, submitter).submit(Long.toString(TABLE_ID), spec()); + verify(actions, never()).tryAcquireInFlightBatch(anyLong(), anyList(), anyLong()); + verify(submitter, never()).runJob(anyString(), anyString(), anyMap()); + } + + @Test + void testSubmitClaimsEveryChainPolicyAndUsesSpecJobConf() { + MaintenanceJobSubmitter submitter = mock(MaintenanceJobSubmitter.class); + when(submitter.runJob(anyString(), anyString(), anyMap())).thenReturn(999L); + TableMaintenanceActionMetaService actions = mock(TableMaintenanceActionMetaService.class); + List<OccupancyClaim> claims = List.of(new OccupancyClaim(1L, 2L), new OccupancyClaim(3L, 4L)); + when(actions.tryAcquireInFlightBatch( + eq(TABLE_ID), eq(List.of(POLICY_ID, EXPIRE_POLICY_ID)), eq(CLAIM_TIMEOUT_MS))) + .thenReturn(Optional.of(claims)); + + Map<String, String> jobConf = new LinkedHashMap<>(); + jobConf.put( + MaintenanceJobTemplates.MAINTENANCE_TYPES_KEY, "rewrite_data_files,expire_snapshots"); + jobConf.put("updater_options", "{\"gravitino_uri\":\"http://localhost:8090\"}";); + + MaintenanceJobSpec spec = + new MaintenanceJobSpec( + List.of(new TableTarget(TABLE_ID, List.of("ml", "iceberg", "db", "t"))), + List.of( + new Operation(POLICY_ID, "rewrite_data_files"), + new Operation(EXPIRE_POLICY_ID, "expire_snapshots")), + jobConf); + + pipeline(allowAllGate(), actions, submitter).submit(Long.toString(TABLE_ID), spec); + + @SuppressWarnings("unchecked") + ArgumentCaptor<Map<String, String>> confCaptor = ArgumentCaptor.forClass(Map.class); + verify(submitter) + .runJob( + eq("ml"), + eq(MaintenanceJobTemplates.TMS_ICEBERG_MAINTENANCE_TEMPLATE), + confCaptor.capture()); + verify(actions).bindJobIds(claims, 999L); + verify(actions, never()).bindJobId(org.mockito.ArgumentMatchers.any(), anyLong()); + + Map<String, String> conf = confCaptor.getValue(); + Assertions.assertEquals( + "rewrite_data_files,expire_snapshots", + conf.get(MaintenanceJobTemplates.MAINTENANCE_TYPES_KEY)); + Assertions.assertFalse(conf.containsKey("tms_table_id")); + Assertions.assertFalse(conf.containsKey("tms_policy_id")); + Assertions.assertEquals("iceberg", conf.get("catalog_name")); + Assertions.assertNull(conf.get("table_identifier")); + Assertions.assertTrue(conf.get("tables").contains("\"table_identifier\":\"db.t\"")); + Assertions.assertTrue(conf.get("tables").contains("\"table_id\":" + TABLE_ID)); + } + + private static MaintenanceCommitSubmitPipeline pipeline( + CommitOccupancyGate gate, + TableMaintenanceActionMetaService actions, + MaintenanceJobSubmitter submitter) { + return new MaintenanceCommitSubmitPipeline( + gate, + actions, + submitter, + new MaintenanceJobConfBuilder(new Config(false) {}), + CLAIM_TIMEOUT_MS); + } + + private static CommitOccupancyGate allowAllGate() { + return new CommitOccupancyGate() { + @Override + public boolean isInFlight(long tableId, long policyId) { + return false; + } + + @Override + public boolean withinMinInterval( + long tableId, long policyId, long minIntervalMs, long nowMs) { + return false; + } + }; + } + + private static CommitOccupancyGate inFlightGate(long... policyIds) { + Set<Long> blocked = new HashSet<>(); + for (long id : policyIds) { + blocked.add(id); + } + return new CommitOccupancyGate() { + @Override + public boolean isInFlight(long tableId, long policyId) { + return blocked.contains(policyId); + } + + @Override + public boolean withinMinInterval( + long tableId, long policyId, long minIntervalMs, long nowMs) { + return false; + } + }; + } + + private static MaintenanceJobSpec spec() { + Map<String, String> jobConf = new LinkedHashMap<>(); + jobConf.put(MaintenanceJobTemplates.MAINTENANCE_TYPES_KEY, "rewrite_data_files"); + jobConf.put("updater_options", "{\"gravitino_uri\":\"http://localhost:8090\"}";); + return new MaintenanceJobSpec( + List.of(new TableTarget(TABLE_ID, List.of("ml", "iceberg", "db", "t"))), + List.of(new Operation(POLICY_ID, "rewrite_data_files")), + jobConf); + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/commit/TestRelationalCommitJobScheduler.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/commit/TestRelationalCommitJobScheduler.java new file mode 100644 index 0000000000..6c6797c898 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/job/commit/TestRelationalCommitJobScheduler.java @@ -0,0 +1,190 @@ +/* + * 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.commit; + +import com.github.kagkarlsson.scheduler.SchedulerClient; +import com.github.kagkarlsson.scheduler.task.TaskInstanceId; +import java.time.Instant; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.function.Function; +import javax.sql.DataSource; +import org.apache.gravitino.Config; +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.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; + +class TestRelationalCommitJobScheduler extends TestJDBCBackend { + + private static final long TABLE_ID = 101L; + private static final long POLICY_ID = 202L; + private static final String GRAVITINO_URI = "http://localhost:8090";; + + @TestTemplate + void testEnqueueSkipsWhenAllWithinMinInterval() { + RelationalCommitJobScheduler scheduler = + new RelationalCommitJobScheduler( + dataSource(), + configWithUri(), + fixedCollector(List.of(member(POLICY_ID))), + coolingGate(POLICY_ID)); + Assertions.assertEquals(CommitJobUpsertResult.SKIPPED, scheduler.upsertDueNow(TABLE_ID)); + Assertions.assertTrue( + RelationalCommitJobScheduler.client(dataSource()) + .getScheduledExecution( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.COMMIT_JOB, Long.toString(TABLE_ID))) + .isEmpty()); + } + + @TestTemplate + void testEnqueueWritesSpecAndIsIdempotent() { + RelationalCommitJobScheduler scheduler = + new RelationalCommitJobScheduler( + dataSource(), + configWithUri(), + fixedCollector(List.of(member(POLICY_ID))), + allowAllGate()); + + Assertions.assertEquals(CommitJobUpsertResult.INSERTED, scheduler.upsertDueNow(TABLE_ID)); + Assertions.assertEquals(CommitJobUpsertResult.ALREADY_QUEUED, scheduler.upsertDueNow(TABLE_ID)); + + SchedulerClient client = RelationalCommitJobScheduler.client(dataSource()); + MaintenanceJobSpec spec = + (MaintenanceJobSpec) + client + .getScheduledExecution( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.COMMIT_JOB, Long.toString(TABLE_ID))) + .get() + .getData(); + Assertions.assertTrue(spec.isValid()); + Assertions.assertEquals(List.of(TABLE_ID), spec.tableIds()); + Assertions.assertEquals(POLICY_ID, spec.getOperations().get(0).getPolicyId()); + Assertions.assertTrue(spec.getJobConf().containsKey("updater_options")); + Assertions.assertNull(spec.getJobConf().get("catalog_name")); + Assertions.assertNull(spec.getJobConf().get("table_identifier")); + Assertions.assertEquals("iceberg", spec.getTables().get(0).catalogName()); + } + + private static Function<Long, Optional<CollectedCommitChain>> fixedCollector( + List<CommitChainMember> chain) { + ResolvedTableRef ref = new ResolvedTableRef(TABLE_ID, "iceberg", "db.t"); + CollectedCommitChain collected = new CollectedCommitChain(chain, ref); + return tableId -> tableId == TABLE_ID ? Optional.of(collected) : Optional.empty(); + } + + private static CommitChainMember member(long policyId) { + return new CommitChainMember( + Policy.BuiltInType.ICEBERG_COMPACTION, policy(policyId), "rewrite_data_files"); + } + + private static CommitOccupancyGate allowAllGate() { + return new CommitOccupancyGate() { + @Override + public boolean isInFlight(long tableId, long policyId) { + return false; + } + + @Override + public boolean withinMinInterval( + long tableId, long policyId, long minIntervalMs, long nowMs) { + return false; + } + }; + } + + private static CommitOccupancyGate coolingGate(long... policyIds) { + Set<Long> cooling = new HashSet<>(); + for (long id : policyIds) { + cooling.add(id); + } + return new CommitOccupancyGate() { + @Override + public boolean isInFlight(long tableId, long policyId) { + return false; + } + + @Override + public boolean withinMinInterval( + long tableId, long policyId, long minIntervalMs, long nowMs) { + return cooling.contains(policyId); + } + }; + } + + 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 PolicyEntity policy(long id) { + return PolicyEntity.builder() + .withId(id) + .withName("p-" + id) + .withNamespace(NamespaceUtil.ofPolicy("ml")) + .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION) + .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("gravitino_uri", GRAVITINO_URI))) + .withEnabled(true) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.EPOCH).build()) + .build(); + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/storage/service/TestTableMaintenanceActionMetaService.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/storage/service/TestTableMaintenanceActionMetaService.java index c5c694ae7b..515b59b6bb 100644 --- a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/storage/service/TestTableMaintenanceActionMetaService.java +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/storage/service/TestTableMaintenanceActionMetaService.java @@ -21,6 +21,7 @@ package org.apache.gravitino.maintenance.storage.service; import static org.apache.gravitino.maintenance.storage.TableMaintenanceActionTestUtils.expireClaim; +import java.util.List; import java.util.Optional; import org.apache.gravitino.maintenance.TableMaintenanceConfig; import org.apache.gravitino.storage.relational.TestJDBCBackend; @@ -42,6 +43,7 @@ class TestTableMaintenanceActionMetaService extends TestJDBCBackend { long after = System.currentTimeMillis(); Assertions.assertNull(actionService.getById(claim.id()).jobId()); Assertions.assertEquals(0L, actionService.getById(claim.id()).finishedAt()); + // DB clock may skew slightly from the JVM; allow a small window. Assertions.assertTrue(claim.createdAt() >= before - 5_000L); Assertions.assertTrue(claim.createdAt() <= after + 5_000L); Assertions.assertEquals(claim.createdAt(), actionService.getById(claim.id()).createdAt()); @@ -153,6 +155,24 @@ class TestTableMaintenanceActionMetaService extends TestJDBCBackend { Assertions.assertEquals(jobId, actions.getByJobIdAndTableId(jobId, tableId).jobId()); } + @TestTemplate + void testBindJobIdsIdempotentWhenAlreadyBoundToSameJob() { + TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); + TableMaintenanceConfig config = TableMaintenanceConfig.forTest(1, 1, 1, 500L, 1_000L, 4, 10); + long tableId = 37L; + long policyA = 48L; + long policyB = 49L; + long jobId = 5700L; + OccupancyClaim claimA = actions.insertInFlight(tableId, policyA, config); + OccupancyClaim claimB = actions.insertInFlight(tableId, policyB, config); + actions.bindJobId(claimA, jobId); + actions.bindJobId(claimB, jobId); + // Spark may lazy-bind before commit-chain submit-side bindJobIds. + actions.bindJobIds(List.of(claimA, claimB), jobId); + Assertions.assertEquals(jobId, actions.getById(claimA.id()).jobId()); + Assertions.assertEquals(jobId, actions.getById(claimB.id()).jobId()); + } + @TestTemplate void testBindJobIdFailsWhenBoundToDifferentJob() { TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); @@ -265,7 +285,7 @@ class TestTableMaintenanceActionMetaService extends TestJDBCBackend { TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); long tableId = 12L; long policyId = 22L; - long timeoutMs = 200L; + long timeoutMs = TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS; Optional<OccupancyClaim> first = actions.tryAcquireInFlight(tableId, policyId, timeoutMs); Assertions.assertTrue(first.isPresent()); @@ -274,11 +294,14 @@ class TestTableMaintenanceActionMetaService extends TestJDBCBackend { Optional<OccupancyClaim> second = actions.tryAcquireInFlight(tableId, policyId, timeoutMs); Assertions.assertTrue(second.isPresent()); - Thread.sleep(timeoutMs + 50L); + OccupancyClaim expired = expireClaim(second.get(), timeoutMs); + // Let the DB clock advance so reclaim cannot reuse the original fencing token. + Thread.sleep(2L); Optional<OccupancyClaim> reclaimed = actions.tryAcquireInFlight(tableId, policyId, timeoutMs); Assertions.assertTrue(reclaimed.isPresent()); Assertions.assertEquals(second.get().id(), reclaimed.get().id()); - Assertions.assertTrue(reclaimed.get().createdAt() >= second.get().createdAt()); + Assertions.assertTrue(reclaimed.get().createdAt() > expired.createdAt()); + Assertions.assertNotEquals(second.get().createdAt(), reclaimed.get().createdAt()); // Stale submitter cleanup must not delete the reclaimed live claim. Assertions.assertFalse(actions.deleteUnboundClaim(second.get())); @@ -289,4 +312,63 @@ class TestTableMaintenanceActionMetaService extends TestJDBCBackend { Assertions.assertTrue(actions.deleteUnboundClaim(reclaimed.get())); Assertions.assertNull(actions.getById(reclaimed.get().id())); } + + @TestTemplate + void testDeleteClaimsFencesAgainstReclaim() throws Exception { + TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); + long tableId = 38L; + long policyId = 50L; + long timeoutMs = TableMaintenanceConfig.DEFAULT_CLAIM_TIMEOUT_MS; + + Optional<OccupancyClaim> first = actions.tryAcquireInFlight(tableId, policyId, timeoutMs); + Assertions.assertTrue(first.isPresent()); + OccupancyClaim expired = expireClaim(first.get(), timeoutMs); + // Let the DB clock advance so reclaim cannot reuse the original fencing token. + Thread.sleep(2L); + Optional<OccupancyClaim> reclaimed = actions.tryAcquireInFlight(tableId, policyId, timeoutMs); + Assertions.assertTrue(reclaimed.isPresent()); + Assertions.assertEquals(first.get().id(), reclaimed.get().id()); + Assertions.assertTrue(reclaimed.get().createdAt() > expired.createdAt()); + + // Failed submit cleanup uses deleteClaims; stale created_at must not remove live claim. + actions.deleteClaims(List.of(first.get())); + Assertions.assertNotNull(actions.getById(reclaimed.get().id())); + Assertions.assertNull(actions.getById(reclaimed.get().id()).jobId()); + Assertions.assertEquals(0L, actions.getById(reclaimed.get().id()).finishedAt()); + + actions.deleteClaims(List.of(reclaimed.get())); + Assertions.assertNull(actions.getById(reclaimed.get().id())); + } + + @TestTemplate + void testTryAcquireBatchClaimsAllPoliciesOrNone() { + TableMaintenanceActionMetaService actions = TableMaintenanceActionMetaService.getInstance(); + long tableId = 31L; + long policyA = 41L; + long policyB = 42L; + long timeoutMs = 60_000L; + + Optional<List<OccupancyClaim>> batch = + actions.tryAcquireInFlightBatch(tableId, List.of(policyA, policyB), timeoutMs); + Assertions.assertTrue(batch.isPresent()); + Assertions.assertEquals(2, batch.get().size()); + Assertions.assertEquals(0L, actions.getById(batch.get().get(0).id()).finishedAt()); + Assertions.assertEquals(0L, actions.getById(batch.get().get(1).id()).finishedAt()); + + // Partial hold blocks the whole chain. + Assertions.assertTrue( + actions.tryAcquireInFlightBatch(tableId, List.of(policyA, policyB), timeoutMs).isEmpty()); + Assertions.assertTrue(actions.tryAcquireInFlight(tableId, policyA, timeoutMs).isEmpty()); + + long jobId = 3100L; + actions.bindJobIds(batch.get(), jobId); + Assertions.assertEquals(jobId, actions.getById(batch.get().get(0).id()).jobId()); + Assertions.assertEquals(jobId, actions.getById(batch.get().get(1).id()).jobId()); + + long finishedAt = System.currentTimeMillis(); + actions.markFinished(jobId, tableId, policyA, 500L, finishedAt); + actions.markFinished(jobId, tableId, policyB, 501L, finishedAt + 1); + Assertions.assertTrue(actions.getById(batch.get().get(0).id()).finishedAt() > 0); + Assertions.assertTrue(actions.getById(batch.get().get(1).id()).finishedAt() > 0); + } }
