diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/CronJobData.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/CronJobData.java new file mode 100644 index 00000000000..0b20c283564 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/CronJobData.java @@ -0,0 +1,72 @@ +/* + * 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; + +import java.io.Serializable; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.Objects; + +/** + * Payload for {@code cron-job} rows. One table for {@code table:...} instance ids; many tables for + * batched snapshot-expiry {@code batch:...} instance ids (§5.2.2). + */ +public final class CronJobData implements Serializable { + + private static final long serialVersionUID = 1L; + + private final ArrayList tableIds; + + /** + * @param tableIds table entity ids covered by this cron-job + */ + public CronJobData(List tableIds) { + this.tableIds = new ArrayList<>(tableIds == null ? List.of() : tableIds); + } + + /** + * @return unmodifiable table ids + */ + public List tableIds() { + return Collections.unmodifiableList(tableIds); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof CronJobData)) { + return false; + } + CronJobData that = (CronJobData) o; + return Objects.equals(tableIds, that.tableIds); + } + + @Override + public int hashCode() { + return Objects.hash(tableIds); + } + + @Override + public String toString() { + return "CronJobData{tableIds=" + tableIds + '}'; + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/CronJobIds.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/CronJobIds.java new file mode 100644 index 00000000000..68e23201411 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/CronJobIds.java @@ -0,0 +1,79 @@ +/* + * 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; + +import java.nio.charset.StandardCharsets; +import java.security.MessageDigest; +import java.security.NoSuchAlgorithmException; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.stream.Collectors; + +/** + * Builds db-scheduler instance ids for one-shot {@code cron-job} rows (§5.2.2). + * + *

Most types use {@code table:{table_id}:{policy_id}}. Batched snapshot-expiry uses {@code + * batch:{batch_key}:{policy_id}} where {@code batch_key} is a stable hash of the table id set. + */ +public final class CronJobIds { + + 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:" + tableId + ":" + policyId; + } + + /** + * @param policyId policy entity id + * @param tableIds tables covered by the batch (order ignored) + * @return cron-job instance id for a snapshot-expiry batch + */ + public static String forBatch(long policyId, List tableIds) { + return "batch:" + batchKey(tableIds) + ":" + policyId; + } + + /** + * Stable short hash of the sorted table-id set. + * + * @param tableIds table entity ids + * @return hex batch key + */ + public static String batchKey(List tableIds) { + List sorted = new ArrayList<>(tableIds == null ? List.of() : tableIds); + Collections.sort(sorted); + String material = sorted.stream().map(String::valueOf).collect(Collectors.joining(",")); + try { + MessageDigest digest = MessageDigest.getInstance("SHA-256"); + byte[] hash = digest.digest(material.getBytes(StandardCharsets.UTF_8)); + StringBuilder hex = new StringBuilder(16); + for (int i = 0; i < 8; i++) { + hex.append(String.format("%02x", hash[i])); + } + return hex.toString(); + } catch (NoSuchAlgorithmException e) { + throw new IllegalStateException("SHA-256 not available", e); + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/CronJobScheduler.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/CronJobScheduler.java new file mode 100644 index 00000000000..335aa20ff8b --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/CronJobScheduler.java @@ -0,0 +1,48 @@ +/* + * 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; + +import java.util.List; +import org.apache.gravitino.policy.Policy; + +/** + * Inserts one-shot {@code cron-job} rows into the cron-pool db-scheduler table. + * + *

Existing instance ids should be left untouched so a repeated expand pick does not reset an + * already-queued submit. Snapshot-expiry winners attached at schema/catalog may be merged into + * {@code batch:{batch_key}:{policy_id}} rows (§5.2.2). + */ +@FunctionalInterface +public interface CronJobScheduler { + + /** + * Enqueues due-now cron jobs for the given expand candidates. + * + * @param policyType maintenance built-in type + * @param policyId policy id encoded in each instance id + * @param candidates enqueueable tables with winner levels + * @param batchSize max tables considered for this expand pick (caller may already truncate) + * @return number of newly scheduled cron-job rows + */ + int enqueueDueNow( + Policy.BuiltInType policyType, + long policyId, + List candidates, + int batchSize); +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/ExpandOccupancyGate.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/ExpandOccupancyGate.java new file mode 100644 index 00000000000..d3df7f748d5 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/ExpandOccupancyGate.java @@ -0,0 +1,37 @@ +/* + * 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; + +/** + * Gate used by expand before inserting a {@code cron-job} row. + * + *

Production wires {@link org.apache.gravitino.maintenance.occupancy.OccupancyGate#shouldSkip}. + */ +@FunctionalInterface +public interface ExpandOccupancyGate { + + /** + * @param tableId table entity id + * @param policyId policy entity id + * @param minIntervalMs resolved cooldown for the policy type + * @param nowMs current epoch millis + * @return true when expand must not enqueue a cron-job for this pair + */ + boolean shouldSkip(long tableId, long policyId, long minIntervalMs, long nowMs); +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/ExpandTableCandidate.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/ExpandTableCandidate.java new file mode 100644 index 00000000000..bcf7368e9a1 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/ExpandTableCandidate.java @@ -0,0 +1,88 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.maintenance.scheduler; + +import com.google.common.base.Preconditions; +import java.util.Objects; +import org.apache.gravitino.MetadataObject; + +/** + * One enqueueable table after nearest-policy and occupancy gating, with the metadata level where + * the winning policy was selected. + */ +public final class ExpandTableCandidate { + + private final long tableId; + private final MetadataObject.Type winnerLevel; + + /** + * @param tableId table entity id + * @param winnerLevel {@link MetadataObject.Type#TABLE}, {@link MetadataObject.Type#SCHEMA}, or + * {@link MetadataObject.Type#CATALOG} + */ + public ExpandTableCandidate(long tableId, MetadataObject.Type winnerLevel) { + Preconditions.checkArgument(tableId > 0, "tableId must be positive"); + Preconditions.checkNotNull(winnerLevel, "winnerLevel must not be null"); + this.tableId = tableId; + this.winnerLevel = winnerLevel; + } + + /** + * @return table entity id + */ + public long tableId() { + return tableId; + } + + /** + * @return level where the nearest winning policy was attached + */ + public MetadataObject.Type winnerLevel() { + return winnerLevel; + } + + /** + * @return true when snapshot-expiry may merge this table into a batch cron-job + */ + public boolean batchableForSnapshotExpiry() { + return winnerLevel == MetadataObject.Type.SCHEMA || winnerLevel == MetadataObject.Type.CATALOG; + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof ExpandTableCandidate)) { + return false; + } + ExpandTableCandidate that = (ExpandTableCandidate) o; + return tableId == that.tableId && winnerLevel == that.winnerLevel; + } + + @Override + public int hashCode() { + return Objects.hash(tableId, winnerLevel); + } + + @Override + public String toString() { + return "ExpandTableCandidate{tableId=" + tableId + ", winnerLevel=" + winnerLevel + '}'; + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/NearestMaintenancePolicyResolver.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/NearestMaintenancePolicyResolver.java new file mode 100644 index 00000000000..8ff1f87c131 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/NearestMaintenancePolicyResolver.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; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.TreeSet; +import java.util.stream.Collectors; +import org.apache.gravitino.Entity; +import org.apache.gravitino.EntityStore; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.NameIdentifier; +import org.apache.gravitino.RelationalEntity; +import org.apache.gravitino.SupportsRelationOperations; +import org.apache.gravitino.exceptions.NoSuchEntityException; +import org.apache.gravitino.json.PolicyAssociationSelectorSerde; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.meta.TagEntity; +import org.apache.gravitino.policy.AllValuesSelector; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.policy.PolicyAssociationSelector; +import org.apache.gravitino.policy.TagValueSelector; +import org.apache.gravitino.tag.TagAssignment; +import org.apache.gravitino.utils.MetadataObjectUtil; +import org.apache.gravitino.utils.NameIdentifierUtil; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Picks the nearest maintenance policy of a given {@link Policy.BuiltInType} for a table. + * + *

Walks the table then its ancestors (nearest first). At each level, policies selected by tags + * directly assigned to that object are considered. The first non-empty level decides: one + * match is the winner; two or more same-type matches are a conflict (fail-closed). + */ +final class NearestMaintenancePolicyResolver { + + private static final Logger LOG = LoggerFactory.getLogger(NearestMaintenancePolicyResolver.class); + + private final EntityStore entityStore; + + NearestMaintenancePolicyResolver(EntityStore entityStore) { + this.entityStore = entityStore; + } + + /** + * @param metalake metalake name + * @param table table metadata object + * @param policyType maintenance built-in type + * @return winner, none, or same-level conflict + */ + NearestPolicyDecision resolve( + String metalake, MetadataObject table, Policy.BuiltInType policyType) { + List chain = new ArrayList<>(); + chain.add(table); + chain.addAll(MetadataObjectUtil.getParentMetadataObjects(table)); + for (MetadataObject object : chain) { + TreeSet atLevel = matchingPolicyIds(metalake, object, policyType); + if (atLevel.isEmpty()) { + continue; + } + if (atLevel.size() > 1) { + return NearestPolicyDecision.conflict(object, new ArrayList<>(atLevel)); + } + return NearestPolicyDecision.winner(atLevel.first(), object); + } + return NearestPolicyDecision.none(); + } + + private TreeSet matchingPolicyIds( + String metalake, MetadataObject object, Policy.BuiltInType policyType) { + NameIdentifier objectIdent = MetadataObjectUtil.toEntityIdent(metalake, object); + Entity.EntityType objectType = MetadataObjectUtil.toEntityType(object); + List tags; + try { + tags = + entityStore + .relationOperations() + .listEntitiesByRelation( + SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL, objectIdent, objectType); + } catch (NoSuchEntityException e) { + return new TreeSet<>(); + } catch (IOException e) { + throw new UncheckedIOException( + "Failed to list tags for metadata object " + object.fullName(), e); + } + if (tags.isEmpty()) { + return new TreeSet<>(); + } + + Map tagsByName = + tags.stream() + .collect( + Collectors.toMap( + TagEntity::name, tag -> tag, (left, right) -> left, LinkedHashMap::new)); + List tagIdents = + tagsByName.keySet().stream() + .map(name -> NameIdentifierUtil.ofTag(metalake, name)) + .collect(Collectors.toList()); + + List> relations; + try { + relations = + entityStore + .relationOperations() + .batchListEntitiesByRelation( + SupportsRelationOperations.Type.POLICY_TAG_REL, tagIdents, Entity.EntityType.TAG); + } catch (IOException e) { + LOG.error("Failed to list policies for tags on {}", object.fullName(), e); + throw new UncheckedIOException(e); + } + + TreeSet matched = new TreeSet<>(); + for (RelationalEntity relation : relations) { + TagEntity tag = tagsByName.get(relation.source().name()); + if (tag == null) { + continue; + } + Object target = relation.targetEntity(); + if (!(target instanceof PolicyEntity)) { + continue; + } + PolicyEntity policy = (PolicyEntity) target; + if (!policy.enabled() || policy.policyType() != policyType) { + continue; + } + PolicyAssociationSelector selector = + relation + .relationValue() + .map(PolicyAssociationSelectorSerde::deserialize) + .orElseGet(AllValuesSelector::get); + TagAssignment assignment = tag.assignment().orElseGet(TagAssignment::noValue); + if (matches(selector, assignment)) { + matched.add(policy.id()); + } + } + return matched; + } + + private static boolean matches(PolicyAssociationSelector selector, TagAssignment assignment) { + if (selector instanceof AllValuesSelector) { + return true; + } + if (selector instanceof TagValueSelector) { + String expectedValue = ((TagValueSelector) selector).value(); + return Arrays.asList(assignment.values()).contains(expectedValue); + } + throw new IllegalArgumentException( + "Unsupported policy association selector type: " + selector.type()); + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/NearestPolicyDecision.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/NearestPolicyDecision.java new file mode 100644 index 00000000000..492964002fc --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/NearestPolicyDecision.java @@ -0,0 +1,90 @@ +/* + * 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; + +import java.util.List; +import java.util.Optional; +import javax.annotation.Nullable; +import org.apache.gravitino.MetadataObject; + +/** Result of nearest maintenance-policy resolution for one table and built-in type. */ +final class NearestPolicyDecision { + + enum Kind { + /** Exactly one matching policy on the nearest non-empty level. */ + WINNER, + /** No matching policy on the inheritance chain. */ + NONE, + /** Two or more same-type policies on the nearest non-empty level. */ + CONFLICT + } + + private final Kind kind; + @Nullable private final Long winnerPolicyId; + @Nullable private final MetadataObject winnerLevel; + @Nullable private final MetadataObject conflictLevel; + private final List conflictingPolicyIds; + + private NearestPolicyDecision( + Kind kind, + @Nullable Long winnerPolicyId, + @Nullable MetadataObject winnerLevel, + @Nullable MetadataObject conflictLevel, + List conflictingPolicyIds) { + this.kind = kind; + this.winnerPolicyId = winnerPolicyId; + this.winnerLevel = winnerLevel; + this.conflictLevel = conflictLevel; + this.conflictingPolicyIds = conflictingPolicyIds; + } + + static NearestPolicyDecision winner(long policyId, MetadataObject level) { + return new NearestPolicyDecision(Kind.WINNER, policyId, level, null, List.of()); + } + + static NearestPolicyDecision none() { + return new NearestPolicyDecision(Kind.NONE, null, null, null, List.of()); + } + + static NearestPolicyDecision conflict(MetadataObject level, List policyIds) { + return new NearestPolicyDecision(Kind.CONFLICT, null, null, level, List.copyOf(policyIds)); + } + + Kind kind() { + return kind; + } + + Optional winnerPolicyId() { + return Optional.ofNullable(winnerPolicyId); + } + + @Nullable + MetadataObject winnerLevel() { + return winnerLevel; + } + + @Nullable + MetadataObject conflictLevel() { + return conflictLevel; + } + + List conflictingPolicyIds() { + return conflictingPolicyIds; + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/NearestPolicyTableEnumerator.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/NearestPolicyTableEnumerator.java new file mode 100644 index 00000000000..4668aed9f68 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/NearestPolicyTableEnumerator.java @@ -0,0 +1,287 @@ +/* + * 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; + +import java.io.IOException; +import java.io.UncheckedIOException; +import java.util.ArrayList; +import java.util.LinkedHashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.function.Supplier; +import org.apache.gravitino.Entity; +import org.apache.gravitino.EntityStore; +import org.apache.gravitino.GravitinoEnv; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.MetadataObjects; +import org.apache.gravitino.NameIdentifier; +import org.apache.gravitino.RelationalEntity; +import org.apache.gravitino.meta.GenericEntity; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.meta.SchemaEntity; +import org.apache.gravitino.meta.TableEntity; +import org.apache.gravitino.meta.TagEntity; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.storage.relational.service.PolicyTagRelService; +import org.apache.gravitino.storage.relational.service.SchemaMetaService; +import org.apache.gravitino.storage.relational.service.TableMetaService; +import org.apache.gravitino.storage.relational.service.TagMetaService; +import org.apache.gravitino.utils.NameIdentifierUtil; +import org.apache.gravitino.utils.NamespaceUtil; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Lists tables for which {@code policy} is the nearest maintenance policy of its type. + * + *

Walks {@code policy → tags → metadata objects}, expands catalog/schema attachments to tables, + * then keeps a table only when {@link NearestMaintenancePolicyResolver} selects this policy id and + * {@link ExpandOccupancyGate} allows enqueue. Scanning continues until {@code limit} enqueueable + * tables are found (cooldown / in-flight winners do not consume the batch budget). + */ +public final class NearestPolicyTableEnumerator implements PolicyTableEnumerator { + + private static final Logger LOG = LoggerFactory.getLogger(NearestPolicyTableEnumerator.class); + + private final PolicyTagRelService policyTagRelService; + private final TagMetaService tagMetaService; + private final SchemaMetaService schemaMetaService; + private final TableMetaService tableMetaService; + private final Supplier entityStoreSupplier; + + /** Uses relational singletons and {@link GravitinoEnv#entityStore()}. */ + public NearestPolicyTableEnumerator() { + this( + PolicyTagRelService.getInstance(), + TagMetaService.getInstance(), + SchemaMetaService.getInstance(), + TableMetaService.getInstance(), + () -> GravitinoEnv.getInstance().entityStore()); + } + + /** + * @param policyTagRelService policy-tag relations + * @param tagMetaService tag metadata + * @param schemaMetaService schema listing + * @param tableMetaService table listing + * @param entityStoreSupplier entity store for nearest-policy resolution + */ + public NearestPolicyTableEnumerator( + PolicyTagRelService policyTagRelService, + TagMetaService tagMetaService, + SchemaMetaService schemaMetaService, + TableMetaService tableMetaService, + Supplier entityStoreSupplier) { + this.policyTagRelService = policyTagRelService; + this.tagMetaService = tagMetaService; + this.schemaMetaService = schemaMetaService; + this.tableMetaService = tableMetaService; + this.entityStoreSupplier = entityStoreSupplier; + } + + @Override + public List listTableIds( + PolicyEntity policy, int limit, ExpandOccupancyGate occupancyGate, long minIntervalMs) { + if (limit <= 0) { + return List.of(); + } + Policy.BuiltInType policyType = policy.policyType(); + String metalake = policy.namespace().level(0); + NameIdentifier policyIdent = NameIdentifierUtil.ofPolicy(metalake, policy.name()); + List> tagEdges = + policyTagRelService.listRelations(List.of(policyIdent), Entity.EntityType.POLICY); + + NearestMaintenancePolicyResolver nearest = + new NearestMaintenancePolicyResolver(entityStoreSupplier.get()); + Set seenTables = new LinkedHashSet<>(); + Map enqueueable = new LinkedHashMap<>(); + long nowMs = System.currentTimeMillis(); + ScanContext ctx = + new ScanContext( + metalake, + policy.id(), + policyType, + nearest, + occupancyGate, + minIntervalMs, + nowMs, + seenTables, + enqueueable, + limit); + + for (RelationalEntity edge : tagEdges) { + Object target = edge.targetEntity(); + if (!(target instanceof TagEntity)) { + continue; + } + TagEntity tag = (TagEntity) target; + List objects; + try { + objects = tagMetaService.listAssociatedMetadataObjectsForTag(tag.nameIdentifier()); + } catch (IOException e) { + throw new UncheckedIOException("Failed to list objects for tag " + tag.nameIdentifier(), e); + } + for (GenericEntity object : objects) { + if (enqueueable.size() >= limit) { + return new ArrayList<>(enqueueable.values()); + } + collectCandidates(object, ctx); + } + } + return new ArrayList<>(enqueueable.values()); + } + + private void collectCandidates(GenericEntity object, ScanContext ctx) { + MetadataObject.Type objectType; + try { + objectType = MetadataObject.Type.valueOf(object.type().name()); + } catch (IllegalArgumentException e) { + return; + } + if (object.name() == null || object.name().isBlank()) { + return; + } + MetadataObject metadataObject; + try { + metadataObject = MetadataObjects.parse(object.name(), objectType); + } catch (RuntimeException e) { + return; + } + + switch (objectType) { + case TABLE: + if (object.id() == null) { + return; + } + considerTable(object.id(), metadataObject, ctx); + break; + case SCHEMA: + expandSchema(metadataObject, ctx); + break; + case CATALOG: + expandCatalog(metadataObject, ctx); + break; + default: + break; + } + } + + private void expandCatalog(MetadataObject catalog, ScanContext ctx) { + List schemas = + schemaMetaService.listSchemasByNamespace( + NamespaceUtil.ofSchema(ctx.metalake, catalog.name())); + for (SchemaEntity schema : schemas) { + if (ctx.enqueueable.size() >= ctx.limit) { + return; + } + MetadataObject schemaObject = + MetadataObjects.of(catalog.name(), schema.name(), MetadataObject.Type.SCHEMA); + expandSchema(schemaObject, ctx); + } + } + + private void expandSchema(MetadataObject schema, ScanContext ctx) { + String catalogName = schema.parent(); + if (catalogName == null) { + return; + } + List tables = + tableMetaService.listTablesByNamespace( + NamespaceUtil.ofTable(ctx.metalake, catalogName, schema.name())); + for (TableEntity table : tables) { + if (ctx.enqueueable.size() >= ctx.limit) { + return; + } + MetadataObject tableObject = + NameIdentifierUtil.toMetadataObject(table.nameIdentifier(), Entity.EntityType.TABLE); + considerTable(table.id(), tableObject, ctx); + } + } + + private static void considerTable(long tableId, MetadataObject tableObject, ScanContext ctx) { + if (tableId <= 0 || !ctx.seenTables.add(tableId)) { + return; + } + NearestPolicyDecision decision = ctx.nearest.resolve(ctx.metalake, tableObject, ctx.policyType); + if (decision.kind() == NearestPolicyDecision.Kind.CONFLICT) { + LOG.error( + "Skip table {} for maintenance type {}: ambiguous policies {} on nearest level {} ({})", + tableObject.fullName(), + ctx.policyType, + decision.conflictingPolicyIds(), + decision.conflictLevel() == null ? "?" : decision.conflictLevel().type(), + decision.conflictLevel() == null ? "?" : decision.conflictLevel().fullName()); + return; + } + if (decision.winnerPolicyId().orElse(-1L) != ctx.policyId) { + return; + } + MetadataObject winnerLevel = decision.winnerLevel(); + if (winnerLevel == null) { + return; + } + if (ctx.occupancyGate.shouldSkip(tableId, ctx.policyId, ctx.minIntervalMs, ctx.nowMs)) { + LOG.debug( + "Skip cron-job enqueue for table {} policy {} (in-flight or within minIntervalMs={})", + tableId, + ctx.policyId, + ctx.minIntervalMs); + return; + } + ctx.enqueueable.put(tableId, new ExpandTableCandidate(tableId, winnerLevel.type())); + } + + private static final class ScanContext { + private final String metalake; + private final long policyId; + private final Policy.BuiltInType policyType; + private final NearestMaintenancePolicyResolver nearest; + private final ExpandOccupancyGate occupancyGate; + private final long minIntervalMs; + private final long nowMs; + private final Set seenTables; + private final Map enqueueable; + private final int limit; + + private ScanContext( + String metalake, + long policyId, + Policy.BuiltInType policyType, + NearestMaintenancePolicyResolver nearest, + ExpandOccupancyGate occupancyGate, + long minIntervalMs, + long nowMs, + Set seenTables, + Map enqueueable, + int limit) { + this.metalake = metalake; + this.policyId = policyId; + this.policyType = policyType; + this.nearest = nearest; + this.occupancyGate = occupancyGate; + this.minIntervalMs = minIntervalMs; + this.nowMs = nowMs; + this.seenTables = seenTables; + this.enqueueable = enqueueable; + this.limit = limit; + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/PolicyExpandCrons.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/PolicyExpandCrons.java new file mode 100644 index 00000000000..3bab2e45dbe --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/PolicyExpandCrons.java @@ -0,0 +1,47 @@ +/* + * 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; + +import com.github.kagkarlsson.scheduler.task.ExecutionComplete; +import com.github.kagkarlsson.scheduler.task.schedule.CronSchedule; +import com.github.kagkarlsson.scheduler.task.schedule.CronStyle; +import com.github.kagkarlsson.scheduler.task.schedule.Schedule; +import java.time.Instant; +import java.time.ZoneId; + +/** Quartz cron helpers for {@code policy-expand} next-due calculation. */ +final class PolicyExpandCrons { + + private PolicyExpandCrons() {} + + /** + * Next fire time for a Quartz cron, strictly after now. + * + * @param timeTrigger six-field Quartz cron expression + * @return next execution instant in the system default zone + */ + static Instant nextExecution(String timeTrigger) { + CronSchedule schedule = new CronSchedule(timeTrigger, ZoneId.systemDefault(), CronStyle.QUARTZ); + Instant next = schedule.getNextExecutionTime(ExecutionComplete.simulatedSuccess(Instant.now())); + if (next == null || Schedule.NEVER.equals(next)) { + throw new IllegalArgumentException("timeTrigger has no next execution: " + timeTrigger); + } + return next; + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/PolicyExpandPipeline.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/PolicyExpandPipeline.java new file mode 100644 index 00000000000..48e253b57c1 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/PolicyExpandPipeline.java @@ -0,0 +1,164 @@ +/* + * 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; + +import java.time.Instant; +import java.util.List; +import java.util.Optional; +import javax.annotation.Nullable; +import org.apache.gravitino.Config; +import org.apache.gravitino.maintenance.policy.MaintenanceMinIntervalConfigs; +import org.apache.gravitino.maintenance.policy.MaintenancePolicyDefinitions; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.policy.ExecutablePolicyContent; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.policy.PolicyContent; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +/** + * Expands a due {@code policy-expand} lease into one-shot {@code cron-job} rows (§5.2.1). + * + *

Does not submit Spark. Returns the next expand due time, or empty when the expand row should + * be removed (missing / disabled / blank trigger / non-maintenance). + * + *

Enumeration fills up to {@code enqueueBatchSize} tables that pass nearest-policy and {@code + * minIntervalMs} / in-flight gating (policy content → server conf → code default). + */ +public final class PolicyExpandPipeline { + + private static final Logger LOG = LoggerFactory.getLogger(PolicyExpandPipeline.class); + + private final PolicyLookup policyLookup; + private final PolicyTableEnumerator tableEnumerator; + private final CronJobScheduler cronJobScheduler; + private final ExpandOccupancyGate occupancyGate; + private final Config config; + private final int enqueueBatchSize; + + /** + * @param policyLookup loads the policy for the expand instance id + * @param tableEnumerator lists attached table ids + * @param cronJobScheduler writes cron-job rows + * @param occupancyGate in-flight / minInterval gate + * @param config gravitino server config for per-type minInterval fallbacks + * @param enqueueBatchSize max cron-jobs per expand pick + */ + public PolicyExpandPipeline( + PolicyLookup policyLookup, + PolicyTableEnumerator tableEnumerator, + CronJobScheduler cronJobScheduler, + ExpandOccupancyGate occupancyGate, + Config config, + int enqueueBatchSize) { + this.policyLookup = policyLookup; + this.tableEnumerator = tableEnumerator; + this.cronJobScheduler = cronJobScheduler; + this.occupancyGate = occupancyGate; + this.config = config; + this.enqueueBatchSize = enqueueBatchSize; + } + + /** + * Runs one expand pick for {@code policyId}. + * + * @param policyId expand task instance id + * @return next due time when the expand row should stay, or empty to delete the row + */ + public Optional expand(long policyId) { + Optional loaded = policyLookup.findById(policyId); + if (loaded.isEmpty()) { + LOG.info("Removing expand row for missing policy {}", policyId); + return Optional.empty(); + } + PolicyEntity policy = loaded.get(); + if (!policy.enabled() || !isMaintenance(policy.content())) { + LOG.info( + "Removing expand row for policy {} (enabled={}, maintenance={})", + policyId, + policy.enabled(), + isMaintenance(policy.content())); + return Optional.empty(); + } + String timeTrigger = timeTrigger(policy.content()); + if (timeTrigger == null || timeTrigger.isBlank()) { + LOG.info("Removing expand row for policy {} with blank timeTrigger", policyId); + return Optional.empty(); + } + + // Resolve next due before any cron-job writes so a bad crontab cannot leave orphans. + final Instant nextDue; + try { + nextDue = PolicyExpandCrons.nextExecution(timeTrigger); + } catch (RuntimeException e) { + LOG.warn( + "Removing expand row for policy {} with invalid timeTrigger {}: {}", + policyId, + timeTrigger, + e.toString()); + return Optional.empty(); + } + + long minIntervalMs = + MaintenanceMinIntervalConfigs.resolve( + config, policy.policyType(), policyMinIntervalMs(policy.content())); + List toEnqueue = + tableEnumerator.listTableIds(policy, enqueueBatchSize, occupancyGate, minIntervalMs); + int scheduled = + cronJobScheduler.enqueueDueNow(policy.policyType(), policyId, toEnqueue, enqueueBatchSize); + LOG.info( + "Expanded policy {}: scheduled {} cron-job rows for {} tables (type={}, minIntervalMs={}); next expand due {}", + policyId, + scheduled, + toEnqueue.size(), + policy.policyType(), + minIntervalMs, + nextDue); + return Optional.of(nextDue); + } + + private static boolean isMaintenance(@Nullable PolicyContent content) { + if (content == null) { + return false; + } + for (Policy.BuiltInType type : Policy.BuiltInType.values()) { + if (MaintenancePolicyDefinitions.isMaintenance(type) + && type.contentClass().isInstance(content)) { + return true; + } + } + return false; + } + + @Nullable + private static String timeTrigger(@Nullable PolicyContent content) { + if (!(content instanceof ExecutablePolicyContent)) { + return null; + } + return ((ExecutablePolicyContent) content).executionSpec().timeTrigger(); + } + + @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/scheduler/PolicyLookup.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/PolicyLookup.java new file mode 100644 index 00000000000..08103298d42 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/PolicyLookup.java @@ -0,0 +1,33 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ +package org.apache.gravitino.maintenance.scheduler; + +import java.util.Optional; +import org.apache.gravitino.meta.PolicyEntity; + +/** Loads a policy entity by its persisted id for expand-pool execution. */ +@FunctionalInterface +public interface PolicyLookup { + + /** + * @param policyId persisted policy id (expand task instance id) + * @return the policy, or empty when missing / soft-deleted + */ + Optional findById(long policyId); +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/PolicyTableEnumerator.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/PolicyTableEnumerator.java new file mode 100644 index 00000000000..4c94a0e7b36 --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/PolicyTableEnumerator.java @@ -0,0 +1,42 @@ +/* + * 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; + +import java.util.List; +import org.apache.gravitino.meta.PolicyEntity; + +/** + * Lists enqueueable tables for a maintenance policy expand pick. + * + *

{@code limit} is the max number of enqueueable tables after nearest-policy and occupancy + * gating. Each candidate carries the winner attachment level for snapshot-expiry batching. + */ +@FunctionalInterface +public interface PolicyTableEnumerator { + + /** + * @param policy loaded policy + * @param limit maximum enqueueable tables to return + * @param occupancyGate skip in-flight / cooldown tables while scanning + * @param minIntervalMs resolved cooldown for {@link ExpandOccupancyGate#shouldSkip} + * @return distinct enqueueable candidates, size at most {@code limit} + */ + List listTableIds( + PolicyEntity policy, int limit, ExpandOccupancyGate occupancyGate, long minIntervalMs); +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/RelationalCronJobScheduler.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/RelationalCronJobScheduler.java new file mode 100644 index 00000000000..eb1b2d848ec --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/RelationalCronJobScheduler.java @@ -0,0 +1,208 @@ +/* + * 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; + +import com.github.kagkarlsson.scheduler.SchedulerClient; +import com.github.kagkarlsson.scheduler.task.TaskInstance; +import java.io.PrintWriter; +import java.lang.reflect.InvocationHandler; +import java.lang.reflect.InvocationTargetException; +import java.lang.reflect.Proxy; +import java.sql.Connection; +import java.sql.SQLException; +import java.sql.SQLFeatureNotSupportedException; +import java.time.Instant; +import java.util.ArrayList; +import java.util.List; +import javax.sql.DataSource; +import org.apache.gravitino.policy.Policy; + +/** + * Schedules {@code cron-job} rows through a db-scheduler client on the entity-store DataSource. + * + *

The entity store disables JDBC autocommit. SchedulerClient does not commit in that mode, so + * connections are wrapped to enable autocommit for each schedule call. + */ +public final class RelationalCronJobScheduler implements CronJobScheduler { + + private final DataSource dataSource; + + /** + * @param dataSource relational JDBC pool shared with the entity store + */ + public RelationalCronJobScheduler(DataSource dataSource) { + this.dataSource = dataSource; + } + + @Override + public int enqueueDueNow( + Policy.BuiltInType policyType, + long policyId, + List candidates, + int batchSize) { + if (candidates == null || candidates.isEmpty()) { + return 0; + } + int limit = Math.min(candidates.size(), Math.max(1, batchSize)); + List slice = candidates.subList(0, limit); + + SchedulerClient client = + SchedulerClient.Builder.create( + autocommitDataSource(dataSource), TableMaintenanceSchedulerPools.cronJobTask()) + .tableName(TableMaintenanceSchedulerPools.CRON_TASKS_TABLE) + .build(); + Instant due = Instant.now(); + + if (policyType == Policy.BuiltInType.ICEBERG_SNAPSHOT_EXPIRATION) { + return enqueueSnapshotExpiry(client, policyId, slice, due); + } + int scheduled = 0; + for (ExpandTableCandidate candidate : slice) { + if (scheduleTable(client, policyId, candidate.tableId(), due)) { + scheduled++; + } + } + return scheduled; + } + + private static int enqueueSnapshotExpiry( + SchedulerClient client, long policyId, List slice, Instant due) { + List batchable = new ArrayList<>(); + int scheduled = 0; + for (ExpandTableCandidate candidate : slice) { + if (candidate.batchableForSnapshotExpiry()) { + batchable.add(candidate.tableId()); + } else if (scheduleTable(client, policyId, candidate.tableId(), due)) { + scheduled++; + } + } + if (!batchable.isEmpty() && scheduleBatch(client, policyId, batchable, due)) { + scheduled++; + } + return scheduled; + } + + private static boolean scheduleTable( + SchedulerClient client, long policyId, long tableId, Instant due) { + String id = CronJobIds.forTable(tableId, policyId); + CronJobData data = new CronJobData(List.of(tableId)); + return client.scheduleIfNotExists( + new TaskInstance<>(TableMaintenanceSchedulerPools.CRON_JOB, id, data), due); + } + + private static boolean scheduleBatch( + SchedulerClient client, long policyId, List tableIds, Instant due) { + String id = CronJobIds.forBatch(policyId, tableIds); + CronJobData data = new CronJobData(tableIds); + return client.scheduleIfNotExists( + new TaskInstance<>(TableMaintenanceSchedulerPools.CRON_JOB, id, data), due); + } + + private static DataSource autocommitDataSource(DataSource delegate) { + return new AutocommitDataSource(delegate); + } + + /** + * Returns connections with autocommit enabled so SchedulerClient persists each schedule without + * an external transaction manager. + */ + private static final class AutocommitDataSource implements DataSource { + + private final DataSource delegate; + + AutocommitDataSource(DataSource delegate) { + this.delegate = delegate; + } + + @Override + public Connection getConnection() throws SQLException { + return withAutocommit(delegate.getConnection()); + } + + @Override + public Connection getConnection(String username, String password) throws SQLException { + return withAutocommit(delegate.getConnection(username, password)); + } + + @Override + public PrintWriter getLogWriter() throws SQLException { + return delegate.getLogWriter(); + } + + @Override + public void setLogWriter(PrintWriter out) throws SQLException { + delegate.setLogWriter(out); + } + + @Override + public void setLoginTimeout(int seconds) throws SQLException { + delegate.setLoginTimeout(seconds); + } + + @Override + public int getLoginTimeout() throws SQLException { + return delegate.getLoginTimeout(); + } + + @Override + public java.util.logging.Logger getParentLogger() throws SQLFeatureNotSupportedException { + return delegate.getParentLogger(); + } + + @Override + public T unwrap(Class iface) throws SQLException { + if (iface.isInstance(this)) { + return iface.cast(this); + } + return delegate.unwrap(iface); + } + + @Override + public boolean isWrapperFor(Class iface) throws SQLException { + return iface.isInstance(this) || delegate.isWrapperFor(iface); + } + + private static Connection withAutocommit(Connection connection) throws SQLException { + Connection proxy = + (Connection) + Proxy.newProxyInstance( + Connection.class.getClassLoader(), + new Class[] {Connection.class}, + suppressingSetAutocommit(connection)); + connection.setAutoCommit(true); + return proxy; + } + + private static InvocationHandler suppressingSetAutocommit(Connection connection) { + return (proxy, method, args) -> { + if ("setAutoCommit".equals(method.getName())) { + return null; + } + if ("getAutoCommit".equals(method.getName())) { + return true; + } + try { + return method.invoke(connection, args); + } catch (InvocationTargetException e) { + throw e.getCause(); + } + }; + } + } +} diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/RelationalPolicyLookup.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/RelationalPolicyLookup.java new file mode 100644 index 00000000000..3eb160b0c7f --- /dev/null +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/scheduler/RelationalPolicyLookup.java @@ -0,0 +1,53 @@ +/* + * 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; + +import java.util.Optional; +import org.apache.gravitino.Namespace; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.storage.relational.mapper.MetalakeMetaMapper; +import org.apache.gravitino.storage.relational.mapper.PolicyMetaMapper; +import org.apache.gravitino.storage.relational.po.MetalakePO; +import org.apache.gravitino.storage.relational.po.PolicyPO; +import org.apache.gravitino.storage.relational.utils.POConverters; +import org.apache.gravitino.storage.relational.utils.SessionUtils; +import org.apache.gravitino.utils.NamespaceUtil; + +/** Loads policies through the relational entity-store mappers. */ +public final class RelationalPolicyLookup implements PolicyLookup { + + @Override + public Optional findById(long policyId) { + PolicyPO policyPO = + SessionUtils.getWithoutCommit( + PolicyMetaMapper.class, mapper -> mapper.selectPolicyByPolicyId(policyId)); + if (policyPO == null) { + return Optional.empty(); + } + MetalakePO metalakePO = + SessionUtils.getWithoutCommit( + MetalakeMetaMapper.class, + mapper -> mapper.selectMetalakeMetaById(policyPO.getMetalakeId())); + if (metalakePO == null) { + return Optional.empty(); + } + Namespace namespace = NamespaceUtil.ofPolicy(metalakePO.getMetalakeName()); + return Optional.of(POConverters.fromPolicyPO(policyPO, namespace)); + } +} 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 e60bbf9d182..415fe404912 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 @@ -23,10 +23,17 @@ import com.github.kagkarlsson.scheduler.task.helper.Tasks; import com.google.common.base.Preconditions; import java.time.Duration; +import java.time.Instant; import java.util.ArrayList; import java.util.List; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicReference; +import java.util.function.Supplier; import javax.sql.DataSource; +import org.apache.gravitino.Config; import org.apache.gravitino.maintenance.TableMaintenanceConfig; +import org.apache.gravitino.maintenance.occupancy.OccupancyGate; +import org.apache.gravitino.maintenance.storage.service.TableMaintenanceActionMetaService; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -37,13 +44,24 @@ *

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). * - *

Task handlers are intentionally minimal in this change; submit and occupancy logic land in - * follow-up PRs. + *

The expand pool fans a due {@code policy-expand} lease into {@code cron-job} rows. The cron + * pool deletes each one-shot row after the pick (Spark submit is still deferred). Commit submit + * remains a no-op until a follow-up PR. */ public class TableMaintenanceSchedulerPools implements AutoCloseable { private static final Logger LOG = LoggerFactory.getLogger(TableMaintenanceSchedulerPools.class); + /** + * Expand pipeline supplier used by {@link #policyExpandTask()}. Tests may replace it; production + * {@link #start} installs the relational implementation. + */ + private static final AtomicReference> EXPAND_PIPELINE = + new AtomicReference<>( + () -> { + throw new IllegalStateException("policy-expand pipeline is not configured"); + }); + /** Expand-pool db-scheduler table. */ public static final String EXPAND_TASKS_TABLE = "scheduled_expand_tasks"; @@ -72,21 +90,31 @@ public class TableMaintenanceSchedulerPools implements AutoCloseable { * * @param dataSource relational JDBC pool shared with the entity store * @param config maintenance scheduler settings + * @param gravitinoConfig server config for {@code minIntervalMs} fallbacks */ - public void start(DataSource dataSource, TableMaintenanceConfig config) { + public void start(DataSource dataSource, TableMaintenanceConfig config, Config gravitinoConfig) { Preconditions.checkState(schedulers.isEmpty(), "scheduler pools already started"); Preconditions.checkNotNull(dataSource, "dataSource must not be null"); + Preconditions.checkNotNull(gravitinoConfig, "gravitinoConfig must not be null"); + + PolicyExpandPipeline pipeline = + new PolicyExpandPipeline( + new RelationalPolicyLookup(), + new NearestPolicyTableEnumerator(), + new RelationalCronJobScheduler(dataSource), + new OccupancyGate( + TableMaintenanceActionMetaService.getInstance(), config.claimTimeoutMs()) + ::shouldSkip, + gravitinoConfig, + config.expandEnqueueBatchSize()); + EXPAND_PIPELINE.set(() -> pipeline); List built = List.of( buildScheduler( - dataSource, - config, - config.expandThreads(), - EXPAND_TASKS_TABLE, - noopTask(POLICY_EXPAND)), + dataSource, config, config.expandThreads(), EXPAND_TASKS_TABLE, policyExpandTask()), buildScheduler( - dataSource, config, config.cronThreads(), CRON_TASKS_TABLE, noopTask(CRON_JOB)), + dataSource, config, config.cronThreads(), CRON_TASKS_TABLE, cronJobTask()), buildScheduler( dataSource, config, @@ -116,6 +144,10 @@ public void stop() { } } schedulers.clear(); + EXPAND_PIPELINE.set( + () -> { + throw new IllegalStateException("policy-expand pipeline is not configured"); + }); LOG.info("Stopped table maintenance scheduler pools"); } @@ -129,17 +161,69 @@ private static Scheduler buildScheduler( TableMaintenanceConfig config, int threads, String tableName, - Task task) { + Task task) { return Scheduler.create(dataSource, task) .threads(threads) .pollingInterval(Duration.ofMillis(config.pollingIntervalMs())) .heartbeatInterval(Duration.ofMillis(config.heartbeatIntervalMs())) .missedHeartbeatsLimit(config.missedHeartbeatsLimit()) .tableName(tableName) + .commitWhenAutocommitDisabled(true) .registerShutdownHook() .build(); } + /** + * Expand-pool task definition used by the running expand scheduler. + * + *

On each due pick the handler expands the policy into {@code cron-job} rows and reschedules + * (or removes) the expand lease. + * + * @return the {@code policy-expand} task + */ + public static Task policyExpandTask() { + return Tasks.custom(POLICY_EXPAND, Void.class) + .execute( + (taskInstance, executionContext) -> { + long policyId = Long.parseLong(taskInstance.getId()); + Optional nextDue = EXPAND_PIPELINE.get().get().expand(policyId); + return (executionComplete, executionOperations) -> { + if (nextDue.isEmpty()) { + executionOperations.remove(); + } else { + executionOperations.reschedule(executionComplete, nextDue.get()); + } + }; + }); + } + + /** + * Cron-pool task definition. Spark submit is still deferred; the completion handler removes the + * one-shot row so a later expand pick can {@code scheduleIfNotExists} again (§5.2.2). + * + * @return the {@code cron-job} task + */ + public static Task cronJobTask() { + return Tasks.custom(CRON_JOB, CronJobData.class) + .execute( + (taskInstance, executionContext) -> { + LOG.debug( + "TMS cron-job callback for instance {} tables {} (submit deferred)", + taskInstance.getId(), + taskInstance.getData() == null ? List.of() : taskInstance.getData().tableIds()); + return (executionComplete, executionOperations) -> executionOperations.remove(); + }); + } + + /** + * Replaces the expand pipeline supplier. Intended for unit tests. + * + * @param pipelineSupplier supplies the pipeline used by {@link #policyExpandTask()} + */ + static void setExpandPipelineForTest(Supplier pipelineSupplier) { + EXPAND_PIPELINE.set(pipelineSupplier); + } + private static Task noopTask(String taskName) { return Tasks.custom(taskName, Void.class) .execute( diff --git a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/feature/TableMaintenanceRESTFeature.java b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/feature/TableMaintenanceRESTFeature.java index ec9b9511bca..f86e2da5e84 100644 --- a/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/feature/TableMaintenanceRESTFeature.java +++ b/plugins/tms/src/main/java/org/apache/gravitino/maintenance/web/rest/feature/TableMaintenanceRESTFeature.java @@ -73,7 +73,7 @@ public boolean configure(FeatureContext context) { DataSource dataSource = resolveEntityStoreDataSource(); TableMaintenanceSchedulerPools pools = new TableMaintenanceSchedulerPools(); - pools.start(dataSource, TableMaintenanceConfig.from(config)); + pools.start(dataSource, TableMaintenanceConfig.from(config), config); if (!SCHEDULER_POOLS.compareAndSet(null, pools)) { pools.stop(); LOG.warn("Table maintenance scheduler pools were started concurrently; discarded duplicate"); diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestCronJobIds.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestCronJobIds.java new file mode 100644 index 00000000000..19955057c82 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestCronJobIds.java @@ -0,0 +1,41 @@ +/* + * 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; + +import java.util.List; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.Test; + +class TestCronJobIds { + + @Test + void testForTable() { + Assertions.assertEquals("table:10:7", CronJobIds.forTable(10L, 7L)); + } + + @Test + void testBatchKeyIgnoresOrder() { + Assertions.assertEquals( + CronJobIds.batchKey(List.of(1L, 2L, 3L)), CronJobIds.batchKey(List.of(3L, 1L, 2L))); + Assertions.assertEquals( + CronJobIds.forBatch(9L, List.of(1L, 2L)), CronJobIds.forBatch(9L, List.of(2L, 1L))); + Assertions.assertTrue(CronJobIds.forBatch(9L, List.of(1L, 2L)).startsWith("batch:")); + Assertions.assertTrue(CronJobIds.forBatch(9L, List.of(1L, 2L)).endsWith(":9")); + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestCronJobScheduler.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestCronJobScheduler.java new file mode 100644 index 00000000000..fe9a4b8db16 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestCronJobScheduler.java @@ -0,0 +1,188 @@ +/* + * 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; + +import com.github.kagkarlsson.scheduler.Scheduler; +import com.github.kagkarlsson.scheduler.SchedulerClient; +import com.github.kagkarlsson.scheduler.task.TaskInstanceId; +import java.time.Duration; +import java.util.List; +import java.util.concurrent.TimeUnit; +import javax.sql.DataSource; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.storage.relational.TestJDBCBackend; +import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.TestTemplate; + +class TestCronJobScheduler extends TestJDBCBackend { + + @TestTemplate + void testEnqueueIsIdempotentAndRespectsBatchSize() { + DataSource dataSource = + SqlSessionFactoryHelper.getInstance() + .getSqlSessionFactory() + .getConfiguration() + .getEnvironment() + .getDataSource(); + CronJobScheduler scheduler = new RelationalCronJobScheduler(dataSource); + + List three = + List.of( + candidate(1L, MetadataObject.Type.TABLE), + candidate(2L, MetadataObject.Type.TABLE), + candidate(3L, MetadataObject.Type.TABLE)); + Assertions.assertEquals( + 2, scheduler.enqueueDueNow(Policy.BuiltInType.ICEBERG_COMPACTION, 7L, three, 2)); + Assertions.assertEquals( + 0, + scheduler.enqueueDueNow( + Policy.BuiltInType.ICEBERG_COMPACTION, + 7L, + List.of( + candidate(1L, MetadataObject.Type.TABLE), candidate(2L, MetadataObject.Type.TABLE)), + 2)); + + SchedulerClient client = + SchedulerClient.Builder.create(dataSource, TableMaintenanceSchedulerPools.cronJobTask()) + .tableName(TableMaintenanceSchedulerPools.CRON_TASKS_TABLE) + .build(); + 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()); + } + + @TestTemplate + void testSnapshotExpiryBatchesSchemaLevelTables() { + DataSource dataSource = + SqlSessionFactoryHelper.getInstance() + .getSqlSessionFactory() + .getConfiguration() + .getEnvironment() + .getDataSource(); + CronJobScheduler scheduler = new RelationalCronJobScheduler(dataSource); + + List schemaLevel = + List.of( + candidate(10L, MetadataObject.Type.SCHEMA), + candidate(11L, MetadataObject.Type.SCHEMA), + candidate(12L, MetadataObject.Type.TABLE)); + Assertions.assertEquals( + 2, + scheduler.enqueueDueNow( + Policy.BuiltInType.ICEBERG_SNAPSHOT_EXPIRATION, 9L, schemaLevel, 10)); + + SchedulerClient client = + SchedulerClient.Builder.create(dataSource, TableMaintenanceSchedulerPools.cronJobTask()) + .tableName(TableMaintenanceSchedulerPools.CRON_TASKS_TABLE) + .build(); + Assertions.assertTrue( + client + .getScheduledExecution( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forTable(12L, 9L))) + .isPresent()); + Assertions.assertTrue( + client + .getScheduledExecution( + TaskInstanceId.of( + TableMaintenanceSchedulerPools.CRON_JOB, + CronJobIds.forBatch(9L, List.of(10L, 11L)))) + .isPresent()); + // Idempotent re-enqueue of the same batch. + Assertions.assertEquals( + 0, + scheduler.enqueueDueNow( + Policy.BuiltInType.ICEBERG_SNAPSHOT_EXPIRATION, 9L, schemaLevel, 10)); + } + + @TestTemplate + void testCronJobTaskRemovesRowAfterPick() throws Exception { + DataSource dataSource = + SqlSessionFactoryHelper.getInstance() + .getSqlSessionFactory() + .getConfiguration() + .getEnvironment() + .getDataSource(); + CronJobScheduler enqueue = new RelationalCronJobScheduler(dataSource); + Assertions.assertEquals( + 1, + enqueue.enqueueDueNow( + Policy.BuiltInType.ICEBERG_COMPACTION, + 9L, + List.of(candidate(99L, MetadataObject.Type.TABLE)), + 1)); + + TaskInstanceId instanceId = + TaskInstanceId.of(TableMaintenanceSchedulerPools.CRON_JOB, CronJobIds.forTable(99L, 9L)); + SchedulerClient client = + SchedulerClient.Builder.create(dataSource, TableMaintenanceSchedulerPools.cronJobTask()) + .tableName(TableMaintenanceSchedulerPools.CRON_TASKS_TABLE) + .build(); + Assertions.assertTrue(client.getScheduledExecution(instanceId).isPresent()); + + Scheduler cronPool = + Scheduler.create(dataSource, TableMaintenanceSchedulerPools.cronJobTask()) + .threads(1) + .pollingInterval(Duration.ofMillis(100)) + .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( + Policy.BuiltInType.ICEBERG_COMPACTION, + 9L, + List.of(candidate(99L, MetadataObject.Type.TABLE)), + 1)); + Assertions.assertTrue(client.getScheduledExecution(instanceId).isPresent()); + } + + private static ExpandTableCandidate candidate(long tableId, MetadataObject.Type level) { + return new ExpandTableCandidate(tableId, level); + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestNearestMaintenancePolicyResolver.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestNearestMaintenancePolicyResolver.java new file mode 100644 index 00000000000..f072dcfe47e --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestNearestMaintenancePolicyResolver.java @@ -0,0 +1,215 @@ +/* + * 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; + +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.time.Instant; +import java.util.List; +import org.apache.gravitino.Entity; +import org.apache.gravitino.EntityStore; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.MetadataObjects; +import org.apache.gravitino.RelationalEntity; +import org.apache.gravitino.SupportsRelationOperations; +import org.apache.gravitino.meta.AuditInfo; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.meta.TagEntity; +import org.apache.gravitino.policy.IcebergDataCompactionContent; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.policy.PolicyContents; +import org.apache.gravitino.tag.TagAssignment; +import org.apache.gravitino.utils.MetadataObjectUtil; +import org.apache.gravitino.utils.NameIdentifierUtil; +import org.apache.gravitino.utils.NamespaceUtil; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class TestNearestMaintenancePolicyResolver { + + private static final String METALAKE = "metalake"; + private static final MetadataObject TABLE = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema", "table"), MetadataObject.Type.TABLE); + private static final MetadataObject SCHEMA = + MetadataObjects.of(java.util.Arrays.asList("catalog", "schema"), MetadataObject.Type.SCHEMA); + private static final MetadataObject CATALOG = + MetadataObjects.of( + java.util.Collections.singletonList("catalog"), MetadataObject.Type.CATALOG); + + private EntityStore entityStore; + private SupportsRelationOperations relationOperations; + private NearestMaintenancePolicyResolver resolver; + + @BeforeEach + void setUp() { + entityStore = mock(EntityStore.class); + relationOperations = mock(SupportsRelationOperations.class); + when(entityStore.relationOperations()).thenReturn(relationOperations); + resolver = new NearestMaintenancePolicyResolver(entityStore); + } + + @Test + void testTableLevelBeatsCatalogLevel() throws Exception { + PolicyEntity tablePolicy = policy(10L, "table-compaction"); + PolicyEntity catalogPolicy = policy(20L, "catalog-compaction"); + TagEntity tableTag = tag(1L, "t-table"); + TagEntity catalogTag = tag(2L, "t-catalog"); + + stubDirectTags(TABLE, List.of(tableTag)); + stubDirectTags(SCHEMA, List.of()); + stubDirectTags(CATALOG, List.of(catalogTag)); + + when(relationOperations.batchListEntitiesByRelation( + eq(SupportsRelationOperations.Type.POLICY_TAG_REL), + eq(List.of(NameIdentifierUtil.ofTag(METALAKE, "t-table"))), + eq(Entity.EntityType.TAG))) + .thenReturn(List.of(relation(tableTag, tablePolicy))); + when(relationOperations.batchListEntitiesByRelation( + eq(SupportsRelationOperations.Type.POLICY_TAG_REL), + eq(List.of(NameIdentifierUtil.ofTag(METALAKE, "t-catalog"))), + eq(Entity.EntityType.TAG))) + .thenReturn(List.of(relation(catalogTag, catalogPolicy))); + + NearestPolicyDecision decision = + resolver.resolve(METALAKE, TABLE, Policy.BuiltInType.ICEBERG_COMPACTION); + Assertions.assertEquals(NearestPolicyDecision.Kind.WINNER, decision.kind()); + Assertions.assertEquals(java.util.Optional.of(10L), decision.winnerPolicyId()); + Assertions.assertEquals(TABLE, decision.winnerLevel()); + } + + @Test + void testFallsBackToAncestorWhenTableHasNoMatch() throws Exception { + PolicyEntity catalogPolicy = policy(20L, "catalog-compaction"); + TagEntity catalogTag = tag(2L, "t-catalog"); + + stubDirectTags(TABLE, List.of()); + stubDirectTags(SCHEMA, List.of()); + stubDirectTags(CATALOG, List.of(catalogTag)); + + when(relationOperations.batchListEntitiesByRelation( + eq(SupportsRelationOperations.Type.POLICY_TAG_REL), + eq(List.of(NameIdentifierUtil.ofTag(METALAKE, "t-catalog"))), + eq(Entity.EntityType.TAG))) + .thenReturn(List.of(relation(catalogTag, catalogPolicy))); + + NearestPolicyDecision decision = + resolver.resolve(METALAKE, TABLE, Policy.BuiltInType.ICEBERG_COMPACTION); + Assertions.assertEquals(NearestPolicyDecision.Kind.WINNER, decision.kind()); + Assertions.assertEquals(java.util.Optional.of(20L), decision.winnerPolicyId()); + Assertions.assertEquals(CATALOG, decision.winnerLevel()); + } + + @Test + void testEmptyWhenNoMatchingType() throws Exception { + stubDirectTags(TABLE, List.of()); + stubDirectTags(SCHEMA, List.of()); + stubDirectTags(CATALOG, List.of()); + + NearestPolicyDecision decision = + resolver.resolve(METALAKE, TABLE, Policy.BuiltInType.ICEBERG_COMPACTION); + Assertions.assertEquals(NearestPolicyDecision.Kind.NONE, decision.kind()); + } + + @Test + void testSameLevelConflict() throws Exception { + PolicyEntity p1 = policy(10L, "compaction-a"); + PolicyEntity p2 = policy(11L, "compaction-b"); + TagEntity tagA = tag(1L, "t-a"); + TagEntity tagB = tag(2L, "t-b"); + + stubDirectTags(TABLE, List.of(tagA, tagB)); + + when(relationOperations.batchListEntitiesByRelation( + eq(SupportsRelationOperations.Type.POLICY_TAG_REL), + eq( + List.of( + NameIdentifierUtil.ofTag(METALAKE, "t-a"), + NameIdentifierUtil.ofTag(METALAKE, "t-b"))), + eq(Entity.EntityType.TAG))) + .thenReturn(List.of(relation(tagA, p1), relation(tagB, p2))); + + NearestPolicyDecision decision = + resolver.resolve(METALAKE, TABLE, Policy.BuiltInType.ICEBERG_COMPACTION); + Assertions.assertEquals(NearestPolicyDecision.Kind.CONFLICT, decision.kind()); + Assertions.assertEquals(TABLE, decision.conflictLevel()); + Assertions.assertEquals(List.of(10L, 11L), decision.conflictingPolicyIds()); + Assertions.assertTrue(decision.winnerPolicyId().isEmpty()); + } + + private void stubDirectTags(MetadataObject object, List tags) throws Exception { + doReturn(tags) + .when(relationOperations) + .listEntitiesByRelation( + eq(SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL), + eq(MetadataObjectUtil.toEntityIdent(METALAKE, object)), + eq(MetadataObjectUtil.toEntityType(object))); + } + + private static RelationalEntity relation(TagEntity tag, PolicyEntity policy) { + return new RelationalEntity<>( + SupportsRelationOperations.Type.POLICY_TAG_REL, + tag.nameIdentifier(), + Entity.EntityType.TAG, + policy, + null); + } + + private static TagEntity tag(long id, String name) { + return TagEntity.builder() + .withId(id) + .withName(name) + .withNamespace(NamespaceUtil.ofTag(METALAKE)) + .withAssignment(TagAssignment.noValue()) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.now()).build()) + .build(); + } + + private static PolicyEntity policy(long id, String name) { + return PolicyEntity.builder() + .withId(id) + .withName(name) + .withNamespace(NamespaceUtil.ofPolicy(METALAKE)) + .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, + IcebergDataCompactionContent.DEFAULT_REWRITE_STRATEGY, + IcebergDataCompactionContent.DEFAULT_SORT_ORDER, + null, + "0 0 2 * * ?", + null, + null, + null, + null)) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.now()).build()) + .build(); + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestNearestPolicyTableEnumerator.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestNearestPolicyTableEnumerator.java new file mode 100644 index 00000000000..f6dce804751 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestNearestPolicyTableEnumerator.java @@ -0,0 +1,399 @@ +/* + * 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; + +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.doReturn; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +import java.time.Instant; +import java.util.List; +import org.apache.gravitino.Entity; +import org.apache.gravitino.EntityStore; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.MetadataObjects; +import org.apache.gravitino.RelationalEntity; +import org.apache.gravitino.SupportsRelationOperations; +import org.apache.gravitino.meta.AuditInfo; +import org.apache.gravitino.meta.GenericEntity; +import org.apache.gravitino.meta.PolicyEntity; +import org.apache.gravitino.meta.TableEntity; +import org.apache.gravitino.meta.TagEntity; +import org.apache.gravitino.policy.IcebergDataCompactionContent; +import org.apache.gravitino.policy.Policy; +import org.apache.gravitino.policy.PolicyContents; +import org.apache.gravitino.storage.relational.service.PolicyTagRelService; +import org.apache.gravitino.storage.relational.service.SchemaMetaService; +import org.apache.gravitino.storage.relational.service.TableMetaService; +import org.apache.gravitino.storage.relational.service.TagMetaService; +import org.apache.gravitino.tag.TagAssignment; +import org.apache.gravitino.utils.MetadataObjectUtil; +import org.apache.gravitino.utils.NameIdentifierUtil; +import org.apache.gravitino.utils.NamespaceUtil; +import org.junit.jupiter.api.Assertions; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +class TestNearestPolicyTableEnumerator { + + private static final String METALAKE = "metalake"; + + private PolicyTagRelService policyTagRelService; + private TagMetaService tagMetaService; + private SchemaMetaService schemaMetaService; + private TableMetaService tableMetaService; + private EntityStore entityStore; + private SupportsRelationOperations relationOperations; + private NearestPolicyTableEnumerator enumerator; + + @BeforeEach + void setUp() { + policyTagRelService = mock(PolicyTagRelService.class); + tagMetaService = mock(TagMetaService.class); + schemaMetaService = mock(SchemaMetaService.class); + tableMetaService = mock(TableMetaService.class); + entityStore = mock(EntityStore.class); + relationOperations = mock(SupportsRelationOperations.class); + when(entityStore.relationOperations()).thenReturn(relationOperations); + enumerator = + new NearestPolicyTableEnumerator( + policyTagRelService, + tagMetaService, + schemaMetaService, + tableMetaService, + () -> entityStore); + } + + @Test + void testDirectTableIncludedWhenNearest() throws Exception { + PolicyEntity policy = policy(10L, "compaction"); + TagEntity tag = tag(1L, "maint"); + when(policyTagRelService.listRelations( + eq(List.of(NameIdentifierUtil.ofPolicy(METALAKE, "compaction"))), + eq(Entity.EntityType.POLICY))) + .thenReturn(List.of(policyTagEdge(policy, tag))); + when(tagMetaService.listAssociatedMetadataObjectsForTag(tag.nameIdentifier())) + .thenReturn( + List.of( + GenericEntity.builder() + .withId(100L) + .withName("catalog.schema.t1") + .withEntityType(Entity.EntityType.TABLE) + .build())); + + // Winning path: tags on table select this policy + MetadataObject table = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema", "t1"), MetadataObject.Type.TABLE); + doReturn(List.of(tag)) + .when(relationOperations) + .listEntitiesByRelation( + eq(SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL), + eq(MetadataObjectUtil.toEntityIdent(METALAKE, table)), + eq(Entity.EntityType.TABLE)); + when(relationOperations.batchListEntitiesByRelation( + eq(SupportsRelationOperations.Type.POLICY_TAG_REL), + eq(List.of(NameIdentifierUtil.ofTag(METALAKE, "maint"))), + eq(Entity.EntityType.TAG))) + .thenReturn(List.of(policyFromTag(tag, policy))); + stubEmptyParents(table); + + List got = enumerator.listTableIds(policy, 10, allowAll(), 0L); + Assertions.assertEquals(1, got.size()); + Assertions.assertEquals(100L, got.get(0).tableId()); + Assertions.assertEquals(MetadataObject.Type.TABLE, got.get(0).winnerLevel()); + } + + @Test + void testSchemaExpansionSkipsOverriddenTable() throws Exception { + PolicyEntity catalogPolicy = policy(10L, "catalog-compaction"); + PolicyEntity tablePolicy = policy(11L, "table-compaction"); + TagEntity catalogTag = tag(1L, "maint"); + TagEntity tableTag = tag(2L, "maint-table"); + + when(policyTagRelService.listRelations( + eq(List.of(NameIdentifierUtil.ofPolicy(METALAKE, "catalog-compaction"))), + eq(Entity.EntityType.POLICY))) + .thenReturn(List.of(policyTagEdge(catalogPolicy, catalogTag))); + when(tagMetaService.listAssociatedMetadataObjectsForTag(catalogTag.nameIdentifier())) + .thenReturn( + List.of( + GenericEntity.builder() + .withId(1L) + .withName("catalog.schema") + .withEntityType(Entity.EntityType.SCHEMA) + .build())); + + TableEntity t1 = + TableEntity.builder() + .withId(100L) + .withName("t1") + .withNamespace(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")) + .withAuditInfo(audit()) + .build(); + TableEntity t2 = + TableEntity.builder() + .withId(101L) + .withName("t2") + .withNamespace(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")) + .withAuditInfo(audit()) + .build(); + when(tableMetaService.listTablesByNamespace( + eq(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")))) + .thenReturn(List.of(t1, t2)); + + MetadataObject table1 = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema", "t1"), MetadataObject.Type.TABLE); + MetadataObject table2 = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema", "t2"), MetadataObject.Type.TABLE); + MetadataObject schema = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema"), MetadataObject.Type.SCHEMA); + MetadataObject catalog = + MetadataObjects.of( + java.util.Collections.singletonList("catalog"), MetadataObject.Type.CATALOG); + + // t1 overridden by table-local policy 11 + doReturn(List.of(tableTag)) + .when(relationOperations) + .listEntitiesByRelation( + eq(SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL), + eq(MetadataObjectUtil.toEntityIdent(METALAKE, table1)), + eq(Entity.EntityType.TABLE)); + when(relationOperations.batchListEntitiesByRelation( + eq(SupportsRelationOperations.Type.POLICY_TAG_REL), + eq(List.of(NameIdentifierUtil.ofTag(METALAKE, "maint-table"))), + eq(Entity.EntityType.TAG))) + .thenReturn(List.of(policyFromTag(tableTag, tablePolicy))); + + // t2 inherits catalog policy via schema-level tag + doReturn(List.of()) + .when(relationOperations) + .listEntitiesByRelation( + eq(SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL), + eq(MetadataObjectUtil.toEntityIdent(METALAKE, table2)), + eq(Entity.EntityType.TABLE)); + doReturn(List.of(catalogTag)) + .when(relationOperations) + .listEntitiesByRelation( + eq(SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL), + eq(MetadataObjectUtil.toEntityIdent(METALAKE, schema)), + eq(Entity.EntityType.SCHEMA)); + when(relationOperations.batchListEntitiesByRelation( + eq(SupportsRelationOperations.Type.POLICY_TAG_REL), + eq(List.of(NameIdentifierUtil.ofTag(METALAKE, "maint"))), + eq(Entity.EntityType.TAG))) + .thenReturn(List.of(policyFromTag(catalogTag, catalogPolicy))); + doReturn(List.of()) + .when(relationOperations) + .listEntitiesByRelation( + eq(SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL), + eq(MetadataObjectUtil.toEntityIdent(METALAKE, catalog)), + eq(Entity.EntityType.CATALOG)); + + List got = enumerator.listTableIds(catalogPolicy, 10, allowAll(), 0L); + Assertions.assertEquals(1, got.size()); + Assertions.assertEquals(101L, got.get(0).tableId()); + Assertions.assertEquals(MetadataObject.Type.SCHEMA, got.get(0).winnerLevel()); + } + + @Test + void testFillsBatchSkippingCooldownWinners() throws Exception { + PolicyEntity policy = policy(10L, "compaction"); + TagEntity tag = tag(1L, "maint"); + when(policyTagRelService.listRelations( + eq(List.of(NameIdentifierUtil.ofPolicy(METALAKE, "compaction"))), + eq(Entity.EntityType.POLICY))) + .thenReturn(List.of(policyTagEdge(policy, tag))); + when(tagMetaService.listAssociatedMetadataObjectsForTag(tag.nameIdentifier())) + .thenReturn( + List.of( + GenericEntity.builder() + .withId(1L) + .withName("catalog.schema") + .withEntityType(Entity.EntityType.SCHEMA) + .build())); + + TableEntity t1 = + TableEntity.builder() + .withId(100L) + .withName("t1") + .withNamespace(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")) + .withAuditInfo(audit()) + .build(); + TableEntity t2 = + TableEntity.builder() + .withId(101L) + .withName("t2") + .withNamespace(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")) + .withAuditInfo(audit()) + .build(); + TableEntity t3 = + TableEntity.builder() + .withId(102L) + .withName("t3") + .withNamespace(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")) + .withAuditInfo(audit()) + .build(); + when(tableMetaService.listTablesByNamespace( + eq(NamespaceUtil.ofTable(METALAKE, "catalog", "schema")))) + .thenReturn(List.of(t1, t2, t3)); + + for (String name : List.of("t1", "t2", "t3")) { + MetadataObject table = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema", name), MetadataObject.Type.TABLE); + doReturn(List.of(tag)) + .when(relationOperations) + .listEntitiesByRelation( + eq(SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL), + eq(MetadataObjectUtil.toEntityIdent(METALAKE, table)), + eq(Entity.EntityType.TABLE)); + stubEmptyParents(table); + } + when(relationOperations.batchListEntitiesByRelation( + eq(SupportsRelationOperations.Type.POLICY_TAG_REL), + eq(List.of(NameIdentifierUtil.ofTag(METALAKE, "maint"))), + eq(Entity.EntityType.TAG))) + .thenReturn(List.of(policyFromTag(tag, policy))); + + // First two winners are cooling down; limit=1 must keep scanning to t3. + ExpandOccupancyGate gate = (tableId, policyId, minIntervalMs, nowMs) -> tableId != 102L; + List got = enumerator.listTableIds(policy, 1, gate, 60_000L); + Assertions.assertEquals(1, got.size()); + Assertions.assertEquals(102L, got.get(0).tableId()); + } + + @Test + void testSameLevelConflictSkipsTable() throws Exception { + PolicyEntity policy = policy(10L, "compaction"); + PolicyEntity other = policy(11L, "compaction-other"); + TagEntity tagA = tag(1L, "t-a"); + TagEntity tagB = tag(2L, "t-b"); + + when(policyTagRelService.listRelations( + eq(List.of(NameIdentifierUtil.ofPolicy(METALAKE, "compaction"))), + eq(Entity.EntityType.POLICY))) + .thenReturn(List.of(policyTagEdge(policy, tagA))); + when(tagMetaService.listAssociatedMetadataObjectsForTag(tagA.nameIdentifier())) + .thenReturn( + List.of( + GenericEntity.builder() + .withId(100L) + .withName("catalog.schema.t1") + .withEntityType(Entity.EntityType.TABLE) + .build())); + + MetadataObject table = + MetadataObjects.of( + java.util.Arrays.asList("catalog", "schema", "t1"), MetadataObject.Type.TABLE); + doReturn(List.of(tagA, tagB)) + .when(relationOperations) + .listEntitiesByRelation( + eq(SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL), + eq(MetadataObjectUtil.toEntityIdent(METALAKE, table)), + eq(Entity.EntityType.TABLE)); + when(relationOperations.batchListEntitiesByRelation( + eq(SupportsRelationOperations.Type.POLICY_TAG_REL), + eq( + List.of( + NameIdentifierUtil.ofTag(METALAKE, "t-a"), + NameIdentifierUtil.ofTag(METALAKE, "t-b"))), + eq(Entity.EntityType.TAG))) + .thenReturn(List.of(policyFromTag(tagA, policy), policyFromTag(tagB, other))); + + Assertions.assertTrue(enumerator.listTableIds(policy, 10, allowAll(), 0L).isEmpty()); + } + + private static ExpandOccupancyGate allowAll() { + return (tableId, policyId, minIntervalMs, nowMs) -> false; + } + + private void stubEmptyParents(MetadataObject table) throws Exception { + for (MetadataObject parent : MetadataObjectUtil.getParentMetadataObjects(table)) { + doReturn(List.of()) + .when(relationOperations) + .listEntitiesByRelation( + eq(SupportsRelationOperations.Type.TAG_METADATA_OBJECT_REL), + eq(MetadataObjectUtil.toEntityIdent(METALAKE, parent)), + eq(MetadataObjectUtil.toEntityType(parent))); + } + } + + private static RelationalEntity policyTagEdge(PolicyEntity policy, TagEntity tag) { + return new RelationalEntity<>( + SupportsRelationOperations.Type.POLICY_TAG_REL, + policy.nameIdentifier(), + Entity.EntityType.POLICY, + tag, + null); + } + + private static RelationalEntity policyFromTag(TagEntity tag, PolicyEntity policy) { + return new RelationalEntity<>( + SupportsRelationOperations.Type.POLICY_TAG_REL, + tag.nameIdentifier(), + Entity.EntityType.TAG, + policy, + null); + } + + private static TagEntity tag(long id, String name) { + return TagEntity.builder() + .withId(id) + .withName(name) + .withNamespace(NamespaceUtil.ofTag(METALAKE)) + .withAssignment(TagAssignment.noValue()) + .withAuditInfo(audit()) + .build(); + } + + private static PolicyEntity policy(long id, String name) { + return PolicyEntity.builder() + .withId(id) + .withName(name) + .withNamespace(NamespaceUtil.ofPolicy(METALAKE)) + .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, + IcebergDataCompactionContent.DEFAULT_REWRITE_STRATEGY, + IcebergDataCompactionContent.DEFAULT_SORT_ORDER, + null, + "0 0 2 * * ?", + null, + null, + null, + null)) + .withAuditInfo(audit()) + .build(); + } + + private static AuditInfo audit() { + return AuditInfo.builder().withCreator("test").withCreateTime(Instant.now()).build(); + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestPolicyExpandPipeline.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestPolicyExpandPipeline.java new file mode 100644 index 00000000000..035855d2042 --- /dev/null +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestPolicyExpandPipeline.java @@ -0,0 +1,238 @@ +/* + * 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; + +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyInt; +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 java.util.List; +import java.util.Optional; +import org.apache.gravitino.Config; +import org.apache.gravitino.MetadataObject; +import org.apache.gravitino.Namespace; +import org.apache.gravitino.maintenance.policy.MaintenanceMinIntervalConfigs; +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 TestPolicyExpandPipeline { + + private static final String CRON = "0 0 2 * * ?"; + + @Test + void testMissingPolicyRemovesExpandRow() { + CronJobScheduler cronJobs = mock(CronJobScheduler.class); + PolicyExpandPipeline pipeline = + pipeline(mockLookupEmpty(1L), mock(PolicyTableEnumerator.class), cronJobs, allowAll()); + + Assertions.assertTrue(pipeline.expand(1L).isEmpty()); + verify(cronJobs, never()).enqueueDueNow(any(), anyLong(), any(), anyInt()); + } + + @Test + void testDisabledPolicyRemovesExpandRow() { + CronJobScheduler cronJobs = mock(CronJobScheduler.class); + PolicyExpandPipeline pipeline = + pipeline( + mockLookup(2L, policy(false, CRON, null)), + mock(PolicyTableEnumerator.class), + cronJobs, + allowAll()); + + Assertions.assertTrue(pipeline.expand(2L).isEmpty()); + verify(cronJobs, never()).enqueueDueNow(any(), anyLong(), any(), anyInt()); + } + + @Test + void testBlankTimeTriggerRemovesExpandRow() { + CronJobScheduler cronJobs = mock(CronJobScheduler.class); + PolicyExpandPipeline pipeline = + pipeline( + mockLookup(3L, policy(true, " ", null)), + mock(PolicyTableEnumerator.class), + cronJobs, + allowAll()); + + Assertions.assertTrue(pipeline.expand(3L).isEmpty()); + verify(cronJobs, never()).enqueueDueNow(any(), anyLong(), any(), anyInt()); + } + + @Test + void testInvalidTimeTriggerRemovesExpandRowWithoutEnqueue() { + PolicyTableEnumerator tables = mock(PolicyTableEnumerator.class); + CronJobScheduler cronJobs = mock(CronJobScheduler.class); + PolicyExpandPipeline pipeline = + pipeline(mockLookup(4L, policy(true, "not-a-cron", null)), tables, cronJobs, allowAll()); + + Assertions.assertTrue(pipeline.expand(4L).isEmpty()); + verify(tables, never()).listTableIds(any(), anyInt(), any(), anyLong()); + verify(cronJobs, never()).enqueueDueNow(any(), anyLong(), any(), anyInt()); + } + + @Test + void testExpandEnqueuesCronJobsAndReturnsNextDue() { + PolicyEntity policy = policy(true, CRON, null); + ExpandOccupancyGate gate = allowAll(); + List candidates = + List.of( + new ExpandTableCandidate(10L, MetadataObject.Type.TABLE), + new ExpandTableCandidate(11L, MetadataObject.Type.TABLE)); + PolicyTableEnumerator tables = mock(PolicyTableEnumerator.class); + when(tables.listTableIds( + eq(policy), + eq(2), + eq(gate), + eq(MaintenanceMinIntervalConfigs.DEFAULT_HOURLY_MIN_INTERVAL_MS))) + .thenReturn(candidates); + CronJobScheduler cronJobs = mock(CronJobScheduler.class); + when(cronJobs.enqueueDueNow( + eq(Policy.BuiltInType.ICEBERG_COMPACTION), eq(42L), eq(candidates), eq(2))) + .thenReturn(2); + PolicyExpandPipeline expandPipeline = + pipeline(mockLookup(42L, policy), tables, cronJobs, gate, 2); + + Optional next = expandPipeline.expand(42L); + Assertions.assertTrue(next.isPresent()); + Assertions.assertTrue(next.get().isAfter(Instant.now().minusSeconds(1))); + verify(cronJobs).enqueueDueNow(Policy.BuiltInType.ICEBERG_COMPACTION, 42L, candidates, 2); + } + + @Test + void testPassesResolvedMinIntervalToEnumerator() { + PolicyEntity policy = policy(true, CRON, 60_000L); + ExpandOccupancyGate gate = allowAll(); + List candidates = + List.of(new ExpandTableCandidate(11L, MetadataObject.Type.TABLE)); + PolicyTableEnumerator tables = mock(PolicyTableEnumerator.class); + when(tables.listTableIds(eq(policy), eq(10), eq(gate), eq(60_000L))).thenReturn(candidates); + CronJobScheduler cronJobs = mock(CronJobScheduler.class); + when(cronJobs.enqueueDueNow( + eq(Policy.BuiltInType.ICEBERG_COMPACTION), eq(42L), eq(candidates), eq(10))) + .thenReturn(1); + + PolicyExpandPipeline expandPipeline = + pipeline(mockLookup(42L, policy), tables, cronJobs, gate, 10); + Assertions.assertTrue(expandPipeline.expand(42L).isPresent()); + verify(tables).listTableIds(eq(policy), eq(10), eq(gate), eq(60_000L)); + verify(cronJobs).enqueueDueNow(Policy.BuiltInType.ICEBERG_COMPACTION, 42L, candidates, 10); + } + + @Test + void testUsesServerDefaultMinIntervalWhenPolicyOmits() { + PolicyEntity policy = policy(true, CRON, null); + ExpandOccupancyGate gate = allowAll(); + List candidates = + List.of(new ExpandTableCandidate(10L, MetadataObject.Type.TABLE)); + PolicyTableEnumerator tables = mock(PolicyTableEnumerator.class); + when(tables.listTableIds( + eq(policy), + eq(10), + eq(gate), + eq(MaintenanceMinIntervalConfigs.DEFAULT_HOURLY_MIN_INTERVAL_MS))) + .thenReturn(candidates); + CronJobScheduler cronJobs = mock(CronJobScheduler.class); + when(cronJobs.enqueueDueNow( + eq(Policy.BuiltInType.ICEBERG_COMPACTION), eq(42L), eq(candidates), eq(10))) + .thenReturn(1); + + PolicyExpandPipeline expandPipeline = + pipeline(mockLookup(42L, policy), tables, cronJobs, gate, 10); + Assertions.assertTrue(expandPipeline.expand(42L).isPresent()); + verify(tables) + .listTableIds( + eq(policy), + eq(10), + eq(gate), + eq(MaintenanceMinIntervalConfigs.DEFAULT_HOURLY_MIN_INTERVAL_MS)); + } + + private static PolicyExpandPipeline pipeline( + PolicyLookup lookup, + PolicyTableEnumerator tables, + CronJobScheduler cronJobs, + ExpandOccupancyGate gate) { + return pipeline(lookup, tables, cronJobs, gate, 100); + } + + private static PolicyExpandPipeline pipeline( + PolicyLookup lookup, + PolicyTableEnumerator tables, + CronJobScheduler cronJobs, + ExpandOccupancyGate gate, + int batchSize) { + return new PolicyExpandPipeline( + lookup, tables, cronJobs, gate, new Config(false) {}, batchSize); + } + + private static ExpandOccupancyGate allowAll() { + return (tableId, policyId, minIntervalMs, nowMs) -> false; + } + + private static PolicyLookup mockLookupEmpty(long policyId) { + PolicyLookup lookup = mock(PolicyLookup.class); + when(lookup.findById(policyId)).thenReturn(Optional.empty()); + return lookup; + } + + private static PolicyLookup mockLookup(long policyId, PolicyEntity policy) { + PolicyLookup lookup = mock(PolicyLookup.class); + when(lookup.findById(policyId)).thenReturn(Optional.of(policy)); + return lookup; + } + + private static PolicyEntity policy(boolean enabled, String timeTrigger, Long minIntervalMs) { + Namespace ns = NamespaceUtil.ofPolicy("metalake"); + return PolicyEntity.builder() + .withId(42L) + .withName("compaction") + .withNamespace(ns) + .withPolicyType(Policy.BuiltInType.ICEBERG_COMPACTION) + .withEnabled(enabled) + .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, + IcebergDataCompactionContent.DEFAULT_REWRITE_STRATEGY, + IcebergDataCompactionContent.DEFAULT_SORT_ORDER, + null, + timeTrigger, + null, + minIntervalMs, + null, + null)) + .withAuditInfo( + AuditInfo.builder().withCreator("test").withCreateTime(Instant.now()).build()) + .build(); + } +} diff --git a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestTableMaintenanceSchedulerPools.java b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestTableMaintenanceSchedulerPools.java index 5de4f541b13..916e96a769a 100644 --- a/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestTableMaintenanceSchedulerPools.java +++ b/plugins/tms/src/test/java/org/apache/gravitino/maintenance/scheduler/TestTableMaintenanceSchedulerPools.java @@ -20,6 +20,7 @@ package org.apache.gravitino.maintenance.scheduler; import javax.sql.DataSource; +import org.apache.gravitino.Config; import org.apache.gravitino.maintenance.TableMaintenanceConfig; import org.apache.gravitino.storage.relational.TestJDBCBackend; import org.apache.gravitino.storage.relational.session.SqlSessionFactoryHelper; @@ -47,9 +48,13 @@ void testThreePoolsStartAndStop() { .getEnvironment() .getDataSource(); pools = new TableMaintenanceSchedulerPools(); + Config gravitinoConfig = new Config(false) {}; Assertions.assertDoesNotThrow( () -> - pools.start(dataSource, TableMaintenanceConfig.forTest(1, 1, 1, 500L, 1_000L, 4, 10))); + pools.start( + dataSource, + TableMaintenanceConfig.forTest(1, 1, 1, 500L, 1_000L, 4, 10), + gravitinoConfig)); Assertions.assertDoesNotThrow(pools::stop); } }