gavinchou commented on code in PR #67820: URL: https://github.com/apache/doris/pull/67820#discussion_r3991506433
########## fe/fe-core/src/main/java/org/apache/doris/tso/TSOTransactionTracker.java: ########## @@ -0,0 +1,408 @@ +// 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.doris.tso; + +import org.apache.doris.cloud.proto.Cloud.GetTsoRecoveryTransactionsResponse; +import org.apache.doris.cloud.proto.Cloud.TxnInfoPB; +import org.apache.doris.cloud.proto.Cloud.TxnStatusPB; +import org.apache.doris.common.Pair; +import org.apache.doris.common.UserException; +import org.apache.doris.transaction.GlobalTransactionMgrIface; +import org.apache.doris.transaction.TransactionState; +import org.apache.doris.transaction.TransactionStatus; + +import com.google.common.base.Preconditions; +import com.google.protobuf.ByteString; +import org.apache.logging.log4j.LogManager; +import org.apache.logging.log4j.Logger; + +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Set; +import java.util.TreeMap; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.locks.Condition; +import java.util.concurrent.locks.ReentrantLock; + +/** In-memory commit registrations. Uses the allocator's lock so no allocated TSO can be missed. */ +final class TSOTransactionTracker { + private static final Logger LOG = LogManager.getLogger(TSOTransactionTracker.class); + private static final int CHECK_BATCH_SIZE = 64; + private static final long CHECK_AGE_NANOS = TimeUnit.SECONDS.toNanos(1); + private final ReentrantLock lock; + private final Condition transactionsChanged; + private final Map<Pair<Long, Long>, PendingTransaction> pendingByTxn = new HashMap<>(); + private final TreeMap<Long, PendingTransaction> pendingByTso = new TreeMap<>(); + // Recovery entries stay separate so the durable committed prefix remains frozen until all finish. + private final TreeMap<Long, PendingTransaction> recoveryByTxn = new TreeMap<>(); + private long generation; + private long recoveryFenceTso; + private long recoveryWatermark; + private boolean recoveryReady; + private boolean recoveryLoaded; + private ByteString recoveryStartKey = ByteString.EMPTY; + private long recoveryPollCursor; + private long pollCursor; + + enum WaitResult { + FINISHED, TIMED_OUT, RECOVERING + } + + private static final class PendingTransaction { + private final Pair<Long, Long> identity; + private final long tso; + private final long registeredAtNanos; + private final Set<Long> tableIds; + + private PendingTransaction(Pair<Long, Long> identity, long tso, long nowNanos, Collection<Long> tableIds) { + this.identity = identity; + this.tso = tso; + this.registeredAtNanos = nowNanos; + this.tableIds = new HashSet<>(tableIds); + } + } + + TSOTransactionTracker(ReentrantLock lock) { + this.lock = lock; + this.transactionsChanged = lock.newCondition(); + } + + void reset(long fenceTso) { + Preconditions.checkState(lock.isHeldByCurrentThread()); + Preconditions.checkArgument(fenceTso > 0, "recovery fence TSO must be positive"); + clearForNewGeneration(); + recoveryFenceTso = fenceTso; + } + + void invalidate() { + Preconditions.checkState(lock.isHeldByCurrentThread()); + clearForNewGeneration(); + recoveryFenceTso = 0; + } + + private void clearForNewGeneration() { + generation++; + pendingByTxn.clear(); + pendingByTso.clear(); + recoveryByTxn.clear(); + recoveryWatermark = 0; + recoveryReady = false; + recoveryLoaded = false; + recoveryStartKey = ByteString.EMPTY; + recoveryPollCursor = 0; + pollCursor = 0; + transactionsChanged.signalAll(); + } + + void register(Pair<Long, Long> identity, long tso, long nowNanos, Set<Long> tableIds) { + Preconditions.checkState(lock.isHeldByCurrentThread()); + Preconditions.checkArgument(!tableIds.isEmpty(), "commit registration requires table IDs"); + PendingTransaction recovered = recoveryByTxn.get(identity.second); + if (recovered != null) { + recovered.tableIds.addAll(tableIds); + } + PendingTransaction existing = pendingByTxn.get(identity); + if (existing != null) { + // A timed-out request can still commit using the earlier TSO. + existing.tableIds.addAll(tableIds); + return; + } + PendingTransaction pending = new PendingTransaction(identity, tso, nowNanos, tableIds); + pendingByTxn.put(identity, pending); + Preconditions.checkState(pendingByTso.put(tso, pending) == null); + } + + void replaceFenced(Pair<Long, Long> identity, long rejectedTso, long fenceTso, long newTso, + long nowNanos, Set<Long> tableIds) { + Preconditions.checkState(lock.isHeldByCurrentThread()); + Preconditions.checkArgument(!tableIds.isEmpty(), "commit registration requires table IDs"); + Preconditions.checkArgument(rejectedTso <= fenceTso, "rejected TSO must be fenced"); + Preconditions.checkArgument(newTso > fenceTso, "replacement TSO must be above the fence"); + PendingTransaction existing = pendingByTxn.get(identity); + Set<Long> mergedTableIds = new HashSet<>(); + if (existing != null) { + Preconditions.checkState(existing.tso <= fenceTso, + "registered transaction TSO must be fenced"); + Preconditions.checkState(pendingByTso.remove(existing.tso) == existing); + mergedTableIds.addAll(existing.tableIds); + } + mergedTableIds.addAll(tableIds); + PendingTransaction replacement = new PendingTransaction(identity, newTso, nowNanos, mergedTableIds); + pendingByTxn.put(identity, replacement); + Preconditions.checkState(pendingByTso.put(newTso, replacement) == null); + PendingTransaction recovered = recoveryByTxn.get(identity.second); + if (recovered != null) { + recovered.tableIds.addAll(tableIds); + } + transactionsChanged.signalAll(); + } + + /** Called with the allocator lock after validating endTso against its current clock. */ + WaitResult awaitTransactions(Map<Long, List<Long>> dbToTableIds, long endTso, long remainingNanos) + throws InterruptedException { + Preconditions.checkState(lock.isHeldByCurrentThread()); + if (!recoveryLoaded) { + return WaitResult.RECOVERING; + } + long waitStartNanos = System.nanoTime(); + long waitGeneration = generation; + List<PendingTransaction> remaining = new ArrayList<>(); + for (PendingTransaction pending : recoveryByTxn.values()) { + List<Long> tables = dbToTableIds.get(pending.identity.first); + if (pending.tso <= endTso && tables != null && !Collections.disjoint(tables, pending.tableIds)) { + remaining.add(pending); + } + } + for (PendingTransaction pending : pendingByTso.headMap(endTso, true).values()) { + List<Long> tables = dbToTableIds.get(pending.identity.first); + if (tables != null && !Collections.disjoint(tables, pending.tableIds)) { + remaining.add(pending); + } + } + // Allocation/registration and this snapshot share the lock. Later allocations are outside + // the validated window; only this fixed set can affect the read. awaitNanos releases the lock. + while (true) { + if (generation != waitGeneration) { + return WaitResult.RECOVERING; + } + remaining.removeIf(pending -> pendingByTxn.get(pending.identity) != pending + && recoveryByTxn.get(pending.identity.second) != pending); + if (remaining.isEmpty()) { + return WaitResult.FINISHED; + } + long nanosLeft = remainingNanos - (System.nanoTime() - waitStartNanos); + if (nanosLeft <= 0) { + return WaitResult.TIMED_OUT; + } + transactionsChanged.awaitNanos(nanosLeft); + } + } + + long candidateCommittedTso(long currentTso, long durableCommittedTso) { + Preconditions.checkState(lock.isHeldByCurrentThread()); + if (!recoveryReady) { + return durableCommittedTso; + } + long candidate = pendingByTso.isEmpty() ? currentTso + : Math.min(currentTso, pendingByTso.firstKey() - 1); + Preconditions.checkState(candidate >= durableCommittedTso, "committed TSO must not regress"); + return candidate; + } + + void transactionFinished(long dbId, long txnId) { Review Comment: naming `markTxnFinised` -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
