io.trino.transaction.InMemoryTransactionManager Maven / Gradle / Ivy

Go to download
/*
 * Licensed 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 io.trino.transaction;

import com.google.common.collect.ImmutableSet;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.errorprone.annotations.ThreadSafe;
import com.google.errorprone.annotations.concurrent.GuardedBy;
import io.airlift.concurrent.BoundedExecutor;
import io.airlift.log.Logger;
import io.airlift.units.Duration;
import io.trino.NotInTransactionException;
import io.trino.metadata.Catalog;
import io.trino.metadata.CatalogInfo;
import io.trino.metadata.CatalogManager;
import io.trino.metadata.CatalogMetadata;
import io.trino.spi.TrinoException;
import io.trino.spi.catalog.CatalogName;
import io.trino.spi.connector.CatalogHandle;
import io.trino.spi.connector.ConnectorTransactionHandle;
import io.trino.spi.transaction.IsolationLevel;
import org.joda.time.DateTime;

import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Executor;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Supplier;

import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.ImmutableList.toImmutableList;
import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static com.google.common.util.concurrent.Futures.immediateFailedFuture;
import static com.google.common.util.concurrent.Futures.immediateFuture;
import static com.google.common.util.concurrent.Futures.immediateVoidFuture;
import static com.google.common.util.concurrent.Futures.nonCancellationPropagating;
import static com.google.common.util.concurrent.MoreExecutors.directExecutor;
import static io.airlift.concurrent.MoreFutures.addExceptionCallback;
import static io.trino.metadata.CatalogManager.NO_CATALOGS;
import static io.trino.spi.StandardErrorCode.ADMINISTRATIVELY_KILLED;
import static io.trino.spi.StandardErrorCode.AUTOCOMMIT_WRITE_CONFLICT;
import static io.trino.spi.StandardErrorCode.CATALOG_NOT_FOUND;
import static io.trino.spi.StandardErrorCode.MULTI_CATALOG_WRITE_CONFLICT;
import static io.trino.spi.StandardErrorCode.READ_ONLY_VIOLATION;
import static io.trino.spi.StandardErrorCode.TRANSACTION_ALREADY_ABORTED;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.stream.Collectors.toList;

@ThreadSafe
public class InMemoryTransactionManager
        implements TransactionManager
{
    private static final Logger log = Logger.get(InMemoryTransactionManager.class);

    private final Duration idleTimeout;
    private final int maxFinishingConcurrency;

    private final ConcurrentMap transactions = new ConcurrentHashMap<>();
    private final CatalogManager catalogManager;
    private final Executor finishingExecutor;

    private InMemoryTransactionManager(Duration idleTimeout, int maxFinishingConcurrency, CatalogManager catalogManager, Executor finishingExecutor)
    {
        this.catalogManager = catalogManager;
        requireNonNull(idleTimeout, "idleTimeout is null");
        checkArgument(maxFinishingConcurrency > 0, "maxFinishingConcurrency must be at least 1");
        requireNonNull(finishingExecutor, "finishingExecutor is null");

        this.idleTimeout = idleTimeout;
        this.maxFinishingConcurrency = maxFinishingConcurrency;
        this.finishingExecutor = finishingExecutor;
    }

    public static TransactionManager create(
            TransactionManagerConfig config,
            ScheduledExecutorService idleCheckExecutor,
            CatalogManager catalogManager,
            Executor finishingExecutor)
    {
        InMemoryTransactionManager transactionManager = new InMemoryTransactionManager(config.getIdleTimeout(), config.getMaxFinishingConcurrency(), catalogManager, finishingExecutor);
        transactionManager.scheduleIdleChecks(config.getIdleCheckInterval(), idleCheckExecutor);
        return transactionManager;
    }

    public static TransactionManager createTestTransactionManager()
    {
        // No idle checks needed
        return new InMemoryTransactionManager(new Duration(1, TimeUnit.DAYS), 1, NO_CATALOGS, directExecutor());
    }

    private void scheduleIdleChecks(Duration idleCheckInterval, ScheduledExecutorService idleCheckExecutor)
    {
        idleCheckExecutor.scheduleWithFixedDelay(() -> {
            try {
                cleanUpExpiredTransactions();
            }
            catch (Throwable t) {
                log.error(t, "Unexpected exception while cleaning up expired transactions");
            }
        }, idleCheckInterval.toMillis(), idleCheckInterval.toMillis(), MILLISECONDS);
    }

    private synchronized void cleanUpExpiredTransactions()
    {
        Iterator> iterator = transactions.entrySet().iterator();
        while (iterator.hasNext()) {
            Entry entry = iterator.next();
            if (entry.getValue().isExpired(idleTimeout)) {
                iterator.remove();
                log.info("Removing expired transaction: %s", entry.getKey());
                entry.getValue().asyncAbort();
            }
        }
    }

    @Override
    public boolean transactionExists(TransactionId transactionId)
    {
        return tryGetTransactionMetadata(transactionId).isPresent();
    }

    @Override
    public TransactionInfo getTransactionInfo(TransactionId transactionId)
    {
        return getTransactionMetadata(transactionId).getTransactionInfo();
    }

    @Override
    public Optional getTransactionInfoIfExist(TransactionId transactionId)
    {
        return tryGetTransactionMetadata(transactionId).map(TransactionMetadata::getTransactionInfo);
    }

    @Override
    public List getAllTransactionInfos()
    {
        return transactions.values().stream()
                .map(TransactionMetadata::getTransactionInfo)
                .collect(toImmutableList());
    }

    @Override
    public Set getTransactionsUsingCatalog(CatalogHandle catalogHandle)
    {
        return transactions.values().stream()
                .filter(transactionMetadata -> transactionMetadata.isUsingCatalog(catalogHandle))
                .map(TransactionMetadata::getTransactionId)
                .collect(toImmutableSet());
    }

    @Override
    public TransactionId beginTransaction(boolean autoCommitContext)
    {
        return beginTransaction(DEFAULT_ISOLATION, DEFAULT_READ_ONLY, autoCommitContext);
    }

    @Override
    public TransactionId beginTransaction(IsolationLevel isolationLevel, boolean readOnly, boolean autoCommitContext)
    {
        TransactionId transactionId = TransactionId.create();
        BoundedExecutor executor = new BoundedExecutor(finishingExecutor, maxFinishingConcurrency);
        TransactionMetadata transactionMetadata = new TransactionMetadata(transactionId, isolationLevel, readOnly, autoCommitContext, catalogManager, executor);
        checkState(transactions.put(transactionId, transactionMetadata) == null, "Duplicate transaction ID: %s", transactionId);
        return transactionId;
    }

    @Override
    public List getCatalogs(TransactionId transactionId)
    {
        return getTransactionMetadata(transactionId).listCatalogs();
    }

    @Override
    public List getActiveCatalogs(TransactionId transactionId)
    {
        return getTransactionMetadata(transactionId).getActiveCatalogs();
    }

    @Override
    public Optional getCatalogHandle(TransactionId transactionId, String catalogName)
    {
        return getTransactionMetadata(transactionId).tryRegisterCatalog(new CatalogName(catalogName));
    }

    @Override
    public Optional getOptionalCatalogMetadata(TransactionId transactionId, String catalogName)
    {
        TransactionMetadata transactionMetadata = getTransactionMetadata(transactionId);
        return transactionMetadata.tryRegisterCatalog(new CatalogName(catalogName))
                .map(transactionMetadata::getTransactionCatalogMetadata);
    }

    @Override
    public CatalogMetadata getCatalogMetadata(TransactionId transactionId, CatalogHandle catalogHandle)
    {
        return getTransactionMetadata(transactionId).getTransactionCatalogMetadata(catalogHandle);
    }

    @Override
    public CatalogMetadata getCatalogMetadataForWrite(TransactionId transactionId, CatalogHandle catalogHandle)
    {
        CatalogMetadata catalogMetadata = getCatalogMetadata(transactionId, catalogHandle);
        checkConnectorWrite(transactionId, catalogHandle);
        return catalogMetadata;
    }

    @Override
    public CatalogMetadata getCatalogMetadataForWrite(TransactionId transactionId, String catalogName)
    {
        TransactionMetadata transactionMetadata = getTransactionMetadata(transactionId);

        CatalogHandle catalogHandle = transactionMetadata.tryRegisterCatalog(new CatalogName(catalogName))
                .orElseThrow(() -> new TrinoException(CATALOG_NOT_FOUND, "Catalog '%s' not found".formatted(catalogName)));

        return getCatalogMetadataForWrite(transactionId, catalogHandle);
    }

    @Override
    public ConnectorTransactionHandle getConnectorTransaction(TransactionId transactionId, String catalogName)
    {
        TransactionMetadata transactionMetadata = getTransactionMetadata(transactionId);

        CatalogHandle catalogHandle = transactionMetadata.tryRegisterCatalog(new CatalogName(catalogName))
                .orElseThrow(() -> new TrinoException(CATALOG_NOT_FOUND, "Catalog '%s' not found".formatted(catalogName)));

        return transactionMetadata.getTransactionCatalogMetadata(catalogHandle).getTransactionHandleFor(catalogHandle);
    }

    @Override
    public ConnectorTransactionHandle getConnectorTransaction(TransactionId transactionId, CatalogHandle catalogHandle)
    {
        return getCatalogMetadata(transactionId, catalogHandle).getTransactionHandleFor(catalogHandle);
    }

    private void checkConnectorWrite(TransactionId transactionId, CatalogHandle catalogHandle)
    {
        getTransactionMetadata(transactionId).checkConnectorWrite(catalogHandle);
    }

    @Override
    public void checkAndSetActive(TransactionId transactionId)
    {
        TransactionMetadata metadata = getTransactionMetadata(transactionId);
        metadata.checkOpenTransaction();
        metadata.setActive();
    }

    @Override
    public void trySetActive(TransactionId transactionId)
    {
        tryGetTransactionMetadata(transactionId).ifPresent(TransactionMetadata::setActive);
    }

    @Override
    public void trySetInactive(TransactionId transactionId)
    {
        tryGetTransactionMetadata(transactionId).ifPresent(TransactionMetadata::setInactive);
    }

    private TransactionMetadata getTransactionMetadata(TransactionId transactionId)
    {
        TransactionMetadata transactionMetadata = transactions.get(transactionId);
        if (transactionMetadata == null) {
            throw new NotInTransactionException(transactionId);
        }
        return transactionMetadata;
    }

    private Optional tryGetTransactionMetadata(TransactionId transactionId)
    {
        return Optional.ofNullable(transactions.get(transactionId));
    }

    private ListenableFuture removeTransactionMetadataAsFuture(TransactionId transactionId)
    {
        TransactionMetadata transactionMetadata = transactions.remove(transactionId);
        if (transactionMetadata == null) {
            return immediateFailedFuture(new NotInTransactionException(transactionId));
        }
        return immediateFuture(transactionMetadata);
    }

    @Override
    public ListenableFuture asyncCommit(TransactionId transactionId)
    {
        return nonCancellationPropagating(Futures.transformAsync(removeTransactionMetadataAsFuture(transactionId), TransactionMetadata::asyncCommit, directExecutor()));
    }

    @Override
    public ListenableFuture asyncAbort(TransactionId transactionId)
    {
        return nonCancellationPropagating(Futures.transformAsync(removeTransactionMetadataAsFuture(transactionId), TransactionMetadata::asyncAbort, directExecutor()));
    }

    @Override
    public void blockCommit(TransactionId transactionId, String reason)
    {
        getTransactionMetadata(transactionId).blockCommit(reason);
    }

    @Override
    public void fail(TransactionId transactionId)
    {
        // Mark the transaction as failed, but don't remove it.
        tryGetTransactionMetadata(transactionId).ifPresent(TransactionMetadata::asyncAbort);
    }

    private static  ListenableFuture asVoid(ListenableFuture future)
    {
        return Futures.transform(future, v -> null, directExecutor());
    }

    @ThreadSafe
    private static class TransactionMetadata
    {
        private final DateTime createTime = DateTime.now();
        private final CatalogManager catalogManager;
        private final TransactionId transactionId;
        private final IsolationLevel isolationLevel;
        private final boolean readOnly;
        private final boolean autoCommitContext;
        private final Executor finishingExecutor;

        private final AtomicReference commitBlocked = new AtomicReference<>();
        private final AtomicReference completedSuccessfully = new AtomicReference<>();
        private final AtomicReference idleStartTime = new AtomicReference<>();

        @GuardedBy("this")
        private final Map> registeredCatalogs = new ConcurrentHashMap<>();
        @GuardedBy("this")
        private final Map activeCatalogs = new ConcurrentHashMap<>();
        @GuardedBy("this")
        private final AtomicReference writtenCatalog = new AtomicReference<>();

        public TransactionMetadata(
                TransactionId transactionId,
                IsolationLevel isolationLevel,
                boolean readOnly,
                boolean autoCommitContext,
                CatalogManager catalogManager,
                Executor finishingExecutor)
        {
            this.transactionId = requireNonNull(transactionId, "transactionId is null");
            this.isolationLevel = requireNonNull(isolationLevel, "isolationLevel is null");
            this.readOnly = readOnly;
            this.autoCommitContext = autoCommitContext;
            this.catalogManager = requireNonNull(catalogManager, "catalogManager is null");
            this.finishingExecutor = requireNonNull(finishingExecutor, "finishingExecutor is null");
        }

        public TransactionId getTransactionId()
        {
            return transactionId;
        }

        public void setActive()
        {
            idleStartTime.set(null);
        }

        public void setInactive()
        {
            idleStartTime.set(System.nanoTime());
        }

        public boolean isExpired(Duration idleTimeout)
        {
            Long idleStartTime = this.idleStartTime.get();
            return idleStartTime != null && Duration.nanosSince(idleStartTime).compareTo(idleTimeout) > 0;
        }

        public void blockCommit(String reason)
        {
            commitBlocked.set(requireNonNull(reason, "reason is null"));
        }

        public void checkOpenTransaction()
        {
            Boolean completedStatus = this.completedSuccessfully.get();
            if (completedStatus != null) {
                if (completedStatus) {
                    // Should not happen normally
                    throw new IllegalStateException("Current transaction already committed");
                }
                throw new TrinoException(TRANSACTION_ALREADY_ABORTED, "Current transaction is aborted, commands ignored until end of transaction block");
            }
        }

        public synchronized boolean isUsingCatalog(CatalogHandle catalogHandle)
        {
            return registeredCatalogs.values().stream()
                    .flatMap(Optional::stream)
                    .map(Catalog::getCatalogHandle)
                    .anyMatch(catalogHandle::equals);
        }

        private synchronized List getActiveCatalogs()
        {
            return activeCatalogs.keySet().stream()
                    .map(CatalogHandle::getCatalogName)
                    .distinct()
                    .map(key -> registeredCatalogs.getOrDefault(key, Optional.empty()))
                    .flatMap(Optional::stream)
                    .map(catalog -> new CatalogInfo(catalog.getCatalogName().toString(), catalog.getCatalogHandle(), catalog.getConnectorName()))
                    .collect(toImmutableList());
        }

        private synchronized List listCatalogs()
        {
            // register all known catalogs - but don't verify so failed catalogs can be listed
            catalogManager.getCatalogNames()
                    .forEach(catalogName -> registeredCatalogs.computeIfAbsent(catalogName, catalogManager::getCatalog));

            return registeredCatalogs.values().stream()
                    .filter(Optional::isPresent)
                    .map(Optional::get)
                    .map(catalog -> new CatalogInfo(catalog.getCatalogName().toString(), catalog.getCatalogHandle(), catalog.getConnectorName()))
                    .collect(toImmutableList());
        }

        private synchronized Optional tryRegisterCatalog(CatalogName catalogName)
        {
            Optional catalog = registeredCatalogs.computeIfAbsent(catalogName, catalogManager::getCatalog);
            catalog.ifPresent(Catalog::verify);
            return catalog.map(Catalog::getCatalogHandle);
        }

        private synchronized CatalogMetadata getTransactionCatalogMetadata(CatalogHandle catalogHandle)
        {
            checkOpenTransaction();

            CatalogMetadata catalogMetadata = activeCatalogs.get(catalogHandle.getRootCatalogHandle());
            if (catalogMetadata == null) {
                // catalog name will not be an internal catalog (e.g., information schema) because internal
                // catalog references can only be generated from the main catalog
                checkArgument(!catalogHandle.getType().isInternal(), "Internal catalog handle not allowed: %s", catalogHandle);
                Catalog catalog = registeredCatalogs.getOrDefault(catalogHandle.getCatalogName(), Optional.empty())
                        .orElseThrow(() -> new IllegalArgumentException("No catalog registered for handle: " + catalogHandle));

                catalogMetadata = catalog.beginTransaction(transactionId, isolationLevel, readOnly, autoCommitContext);

                activeCatalogs.put(catalogHandle, catalogMetadata);
            }
            return catalogMetadata;
        }

        public synchronized void checkConnectorWrite(CatalogHandle catalogHandle)
        {
            checkOpenTransaction();
            CatalogMetadata catalogMetadata = activeCatalogs.get(catalogHandle);
            checkArgument(catalogMetadata != null, "Cannot record write for catalog not part of transaction");
            if (readOnly) {
                throw new TrinoException(READ_ONLY_VIOLATION, "Cannot execute write in a read-only transaction");
            }
            if (!writtenCatalog.compareAndSet(null, catalogHandle) && !writtenCatalog.get().equals(catalogHandle)) {
                CatalogName writtenCatalogName = activeCatalogs.get(writtenCatalog.get()).getCatalogName();
                throw new TrinoException(MULTI_CATALOG_WRITE_CONFLICT, "Multi-catalog writes not supported in a single transaction. Already wrote to catalog " + writtenCatalogName);
            }
            if (catalogMetadata.isSingleStatementWritesOnly() && !autoCommitContext) {
                throw new TrinoException(AUTOCOMMIT_WRITE_CONFLICT, "Catalog only supports writes using autocommit: " + catalogMetadata.getCatalogName());
            }
        }

        public synchronized ListenableFuture asyncCommit()
        {
            if (!completedSuccessfully.compareAndSet(null, true)) {
                if (completedSuccessfully.get()) {
                    // Already done
                    return immediateVoidFuture();
                }
                // Transaction already aborted
                return immediateFailedFuture(new TrinoException(TRANSACTION_ALREADY_ABORTED, "Current transaction has already been aborted"));
            }

            String commitBlockedReason = commitBlocked.get();
            if (commitBlockedReason != null) {
                return Futures.transform(
                        abortInternal(),
                        _ -> {
                            throw new TrinoException(ADMINISTRATIVELY_KILLED, commitBlockedReason);
                        },
                        directExecutor());
            }

            CatalogHandle writeCatalogHandle = this.writtenCatalog.get();
            if (writeCatalogHandle == null) {
                ListenableFuture future = asVoid(Futures.allAsList(activeCatalogs.values().stream()
                        .map(catalog -> Futures.submit(catalog::commit, finishingExecutor))
                        .collect(toList())));
                addExceptionCallback(future, throwable -> {
                    abortInternal();
                    log.error(throwable, "Read-only connector should not throw exception on commit");
                });
                return nonCancellationPropagating(future);
            }

            Supplier> commitReadOnlyConnectors = () -> {
                List> futures = activeCatalogs.entrySet().stream()
                        .filter(entry -> !entry.getKey().equals(writeCatalogHandle))
                        .map(Entry::getValue)
                        .map(transactionMetadata -> Futures.submit(transactionMetadata::commit, finishingExecutor))
                        .collect(toList());
                ListenableFuture future = asVoid(Futures.allAsList(futures));
                addExceptionCallback(future, throwable -> log.error(throwable, "Read-only connector should not throw exception on commit"));
                return future;
            };

            CatalogMetadata writeCatalog = activeCatalogs.get(writeCatalogHandle);
            ListenableFuture commitFuture = Futures.submit(writeCatalog::commit, finishingExecutor);
            ListenableFuture readOnlyCommitFuture = Futures.transformAsync(commitFuture, _ -> commitReadOnlyConnectors.get(), directExecutor());
            addExceptionCallback(readOnlyCommitFuture, this::abortInternal);
            return nonCancellationPropagating(readOnlyCommitFuture);
        }

        public synchronized ListenableFuture asyncAbort()
        {
            if (!completedSuccessfully.compareAndSet(null, false)) {
                if (completedSuccessfully.get()) {
                    // Should not happen normally
                    return immediateFailedFuture(new IllegalStateException("Current transaction already committed"));
                }
                // Already done
                return immediateVoidFuture();
            }
            return abortInternal();
        }

        private synchronized ListenableFuture abortInternal()
        {
            // the callbacks in statement performed on another thread so are safe
            List> futures = activeCatalogs.values().stream()
                    .map(catalog -> Futures.submit(catalog::abort, finishingExecutor))
                    .collect(toList());
            ListenableFuture future = asVoid(Futures.allAsList(futures));
            return nonCancellationPropagating(future);
        }

        public TransactionInfo getTransactionInfo()
        {
            Duration idleTime = Optional.ofNullable(idleStartTime.get())
                    .map(Duration::nanosSince)
                    .orElse(new Duration(0, MILLISECONDS));

            // dereferencing this field is safe because the field is atomic, and activeCatalogs is a concurrent map
            @SuppressWarnings("FieldAccessNotGuarded") Optional writtenCatalogName = Optional.ofNullable(this.writtenCatalog.get())
                    .map(activeCatalogs::get)
                    .map(catalogMetadata -> catalogMetadata.getCatalogName().toString());

            // access here is safe here because the map is concurrent
            @SuppressWarnings("FieldAccessNotGuarded") List catalogNames = activeCatalogs.values().stream()
                    .map(catalogMetadata -> catalogMetadata.getCatalogName().toString())
                    .sorted()
                    .toList();

            return new TransactionInfo(
                    transactionId,
                    isolationLevel,
                    readOnly,
                    autoCommitContext,
                    createTime,
                    idleTime,
                    catalogNames,
                    writtenCatalogName,
                    ImmutableSet.copyOf(activeCatalogs.keySet()));
        }
    }
}