io.trino.connector.CatalogPruneTask Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of trino-main Show documentation
Trino - Engine
There is a newer version: 465
/*
 * 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.connector;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.inject.Inject;
import io.airlift.concurrent.ThreadPoolExecutorMBean;
import io.airlift.discovery.client.ServiceDescriptor;
import io.airlift.discovery.client.ServiceSelector;
import io.airlift.discovery.client.ServiceType;
import io.airlift.http.client.HttpClient;
import io.airlift.http.client.Request;
import io.airlift.http.client.Response;
import io.airlift.http.client.ResponseHandler;
import io.airlift.json.JsonCodec;
import io.airlift.log.Logger;
import io.airlift.node.NodeInfo;
import io.airlift.units.Duration;
import io.trino.metadata.CatalogManager;
import io.trino.metadata.ForNodeManager;
import io.trino.server.InternalCommunicationConfig;
import io.trino.spi.connector.CatalogHandle;
import io.trino.transaction.TransactionManager;
import jakarta.annotation.PostConstruct;
import jakarta.annotation.PreDestroy;
import org.weakref.jmx.Managed;
import org.weakref.jmx.Nested;

import java.net.URI;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.atomic.AtomicBoolean;

import static com.google.common.collect.ImmutableSet.toImmutableSet;
import static com.google.common.net.HttpHeaders.CONTENT_TYPE;
import static com.google.common.net.MediaType.JSON_UTF_8;
import static io.airlift.concurrent.Threads.daemonThreadsNamed;
import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom;
import static io.airlift.http.client.JsonBodyGenerator.jsonBodyGenerator;
import static io.airlift.http.client.Request.Builder.preparePost;
import static io.airlift.json.JsonCodec.listJsonCodec;
import static java.util.Objects.requireNonNull;
import static java.util.concurrent.TimeUnit.MILLISECONDS;

public class CatalogPruneTask
{
    private static final Logger log = Logger.get(CatalogPruneTask.class);
    private static final JsonCodec> CATALOG_HANDLES_CODEC = listJsonCodec(CatalogHandle.class);

    private final TransactionManager transactionManager;
    private final CatalogManager catalogManager;
    private final ConnectorServicesProvider connectorServicesProvider;
    private final NodeInfo nodeInfo;
    private final ServiceSelector selector;
    private final HttpClient httpClient;

    private final ScheduledThreadPoolExecutor executor = new ScheduledThreadPoolExecutor(1, daemonThreadsNamed("catalog-prune"));
    private final ThreadPoolExecutorMBean executorMBean = new ThreadPoolExecutorMBean(executor);

    private final boolean enabled;
    private final Duration updateInterval;
    private final boolean httpsRequired;

    private final AtomicBoolean started = new AtomicBoolean();

    @Inject
    public CatalogPruneTask(
            TransactionManager transactionManager,
            CatalogManager catalogManager,
            ConnectorServicesProvider connectorServicesProvider,
            NodeInfo nodeInfo,
            @ServiceType("trino") ServiceSelector selector,
            @ForNodeManager HttpClient httpClient,
            CatalogPruneTaskConfig catalogPruneTaskConfig,
            InternalCommunicationConfig internalCommunicationConfig)
    {
        this.transactionManager = requireNonNull(transactionManager, "transactionManager is null");
        this.catalogManager = requireNonNull(catalogManager, "catalogManager is null");
        this.connectorServicesProvider = requireNonNull(connectorServicesProvider, "connectorServicesProvider is null");
        this.nodeInfo = requireNonNull(nodeInfo, "nodeInfo is null");
        this.selector = requireNonNull(selector, "selector is null");
        this.httpClient = requireNonNull(httpClient, "httpClient is null");

        this.enabled = catalogPruneTaskConfig.isEnabled();
        updateInterval = catalogPruneTaskConfig.getUpdateInterval();

        this.httpsRequired = internalCommunicationConfig.isHttpsRequired();
    }

    @PostConstruct
    public void start()
    {
        if (enabled && !started.getAndSet(true)) {
            executor.scheduleWithFixedDelay(() -> {
                try {
                    pruneWorkerCatalogs();
                }
                catch (Throwable e) {
                    // ignore to avoid getting unscheduled
                    log.warn(e, "Error pruning catalogs");
                }
            }, updateInterval.toMillis(), updateInterval.toMillis(), MILLISECONDS);
        }
    }

    @PreDestroy
    public void shutdown()
    {
        executor.shutdownNow();
    }

    @Managed
    @Nested
    public ThreadPoolExecutorMBean getExecutor()
    {
        return executorMBean;
    }

    @VisibleForTesting
    public void pruneWorkerCatalogs()
    {
        Set online = selector.selectAllServices().stream()
                .filter(descriptor -> !nodeInfo.getNodeId().equals(descriptor.getNodeId()))
                .collect(toImmutableSet());

        // send message to workers to trigger prune
        List activeCatalogs = getActiveCatalogs();
        pruneWorkerCatalogs(online, activeCatalogs);

        // prune all inactive catalogs - we pass an empty set here because manager always retains active catalogs
        connectorServicesProvider.pruneCatalogs(ImmutableSet.of());
    }

    void pruneWorkerCatalogs(Set online, List activeCatalogs)
    {
        for (ServiceDescriptor service : online) {
            URI uri = getHttpUri(service);
            if (uri == null) {
                continue;
            }
            uri = uriBuilderFrom(uri).appendPath("/v1/task/pruneCatalogs").build();
            Request request = preparePost()
                    .setUri(uri)
                    .addHeader(CONTENT_TYPE, JSON_UTF_8.toString())
                    .setBodyGenerator(jsonBodyGenerator(CATALOG_HANDLES_CODEC, activeCatalogs))
                    .build();
            httpClient.executeAsync(request, new ResponseHandler<>()
            {
                @Override
                public Exception handleException(Request request, Exception exception)
                {
                    log.debug(exception, "Error pruning catalogs on server: %s", request.getUri());
                    return exception;
                }

                @Override
                public Object handle(Request request, Response response)
                {
                    log.debug("Pruned catalogs on server: %s", request.getUri());
                    return null;
                }
            });
        }
    }

    private List getActiveCatalogs()
    {
        ImmutableSet.Builder activeCatalogs = ImmutableSet.builder();
        // all catalogs in an active transaction
        transactionManager.getAllTransactionInfos().forEach(info -> activeCatalogs.addAll(info.getActiveCatalogs()));
        // all catalogs currently associated with a name
        activeCatalogs.addAll(catalogManager.getActiveCatalogs());
        return ImmutableList.copyOf(activeCatalogs.build());
    }

    private URI getHttpUri(ServiceDescriptor descriptor)
    {
        String url = descriptor.getProperties().get(httpsRequired ? "https" : "http");
        if (url != null) {
            return URI.create(url);
        }
        return null;
    }
}