io.kestra.jdbc.runner.JdbcIndexer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of jdbc Show documentation
The modern, scalable orchestrator & scheduler open source platform
The newest version!
package io.kestra.jdbc.runner;

import io.kestra.core.metrics.MetricRegistry;
import io.kestra.core.models.executions.LogEntry;
import io.kestra.core.models.executions.MetricEntry;
import io.kestra.core.queues.QueueFactoryInterface;
import io.kestra.core.queues.QueueInterface;
import io.kestra.core.repositories.LogRepositoryInterface;
import io.kestra.core.repositories.MetricRepositoryInterface;
import io.kestra.core.repositories.SaveRepositoryInterface;
import io.kestra.core.runners.Indexer;
import io.kestra.core.runners.IndexerInterface;
import io.kestra.core.server.ServiceStateChangeEvent;
import io.kestra.core.utils.IdUtils;
import io.kestra.core.utils.ListUtils;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicReference;

import io.micronaut.context.event.ApplicationEventPublisher;
import jakarta.annotation.PreDestroy;
import jakarta.inject.Inject;
import jakarta.inject.Named;
import jakarta.inject.Singleton;
import lombok.extern.slf4j.Slf4j;

/**
 * This class is responsible to batch-indexed asynchronously queue messages.
 * Some queue messages are indexed synchronously via the {@link JdbcQueueIndexer}.
 */
@SuppressWarnings("this-escape")
@Slf4j
@Singleton
@JdbcRunnerEnabled
public class JdbcIndexer implements IndexerInterface {
    private final LogRepositoryInterface logRepository;
    private final JdbcQueue logQueue;

    private final MetricRepositoryInterface metricRepository;
    private final JdbcQueue metricQueue;
    private final MetricRegistry metricRegistry;
    private final List receiveCancellations = new ArrayList<>();

    private final String id = IdUtils.create();
    private final AtomicReference state = new AtomicReference<>();
    private final ApplicationEventPublisher eventPublisher;

    @Inject
    public JdbcIndexer(
        LogRepositoryInterface logRepository,
        @Named(QueueFactoryInterface.WORKERTASKLOG_NAMED) QueueInterface logQueue,
        MetricRepositoryInterface metricRepositor,
        @Named(QueueFactoryInterface.METRIC_QUEUE) QueueInterface metricQueue,
        MetricRegistry metricRegistry,
        ApplicationEventPublisher eventPublisher
    ) {
        this.logRepository = logRepository;
        this.logQueue = (JdbcQueue) logQueue;
        this.metricRepository = metricRepositor;
        this.metricQueue = (JdbcQueue) metricQueue;
        this.metricRegistry = metricRegistry;
        this.eventPublisher = eventPublisher;

        setState(ServiceState.CREATED);
    }

    @Override
    public void run() {
        log.debug("Starting the indexer");
        startQueues();
        setState(ServiceState.RUNNING);
    }

    protected void startQueues() {
        this.sendBatch(logQueue, logRepository);
        this.sendBatch(metricQueue, metricRepository);
    }

    protected  void sendBatch(JdbcQueue queueInterface, SaveRepositoryInterface saveRepositoryInterface) {
        this.receiveCancellations.addFirst(queueInterface.receiveBatch(Indexer.class, eithers -> {
            // first, log all deserialization issues
            eithers.stream().filter(either -> either.isRight()).forEach(either -> log.error("unable to deserialize an item: {}", either.getRight().getMessage()));

            // then index all correctly deserialized items
            List items = eithers.stream().filter(either -> either.isLeft()).map(either -> either.getLeft()).toList();
            if (!ListUtils.isEmpty(items)) {
                String itemClassName = items.getFirst().getClass().getName();
                this.metricRegistry.counter(MetricRegistry.METRIC_INDEXER_REQUEST_COUNT, "type", itemClassName).increment();
                this.metricRegistry.counter(MetricRegistry.METRIC_INDEXER_MESSAGE_IN_COUNT, "type", itemClassName).increment(items.size());

                this.metricRegistry.timer(MetricRegistry.METRIC_INDEXER_REQUEST_DURATION, "type", itemClassName).record(() -> {
                    int saved = saveRepositoryInterface.saveBatch(items);
                    this.metricRegistry.counter(MetricRegistry.METRIC_INDEXER_MESSAGE_OUT_COUNT, "type", itemClassName).increment(saved);
                });
            }
        }));
    }

    private void setState(final ServiceState state) {
        this.state.set(state);
        this.eventPublisher.publishEvent(new ServiceStateChangeEvent(this));
    }

    /** {@inheritDoc} **/
    @Override
    public String getId() {
        return id;
    }
    /** {@inheritDoc} **/
    @Override
    public ServiceType getType() {
        return ServiceType.INDEXER;
    }
    /** {@inheritDoc} **/
    @Override
    public ServiceState getState() {
        return state.get();
    }

    @PreDestroy
    @Override
    public void close() {
        setState(ServiceState.TERMINATING);
        this.receiveCancellations.forEach(Runnable::run);
        try {
            stopQueue();
            setState(ServiceState.TERMINATED_GRACEFULLY);
        } catch (IOException e) {
            log.error("Failed to close the queue", e);
            setState(ServiceState.TERMINATED_FORCED);
        }
    }

    protected void stopQueue() throws IOException {
        this.logQueue.close();
        this.metricQueue.close();
    }
}