<commit_before>/*
 * Grakn - A Distributed Semantic Database
 * Copyright (C) 2016  Grakn Labs Limited
 *
 * Grakn is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Grakn is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Grakn. If not, see <http://www.gnu.org/licenses/gpl.txt>.
 *
 */

package ai.grakn.engine.tasks.manager.singlequeue;

import ai.grakn.engine.GraknEngineConfig;
import ai.grakn.engine.TaskId;
import ai.grakn.engine.tasks.BackgroundTask;
import ai.grakn.engine.tasks.ExternalOffsetStorage;
import ai.grakn.engine.tasks.TaskCheckpoint;
import ai.grakn.engine.tasks.TaskConfiguration;
import ai.grakn.engine.tasks.TaskState;
import ai.grakn.engine.tasks.TaskStateStorage;
import ai.grakn.engine.tasks.connection.RedisConnection;
import ai.grakn.engine.util.EngineID;
import com.codahale.metrics.Meter;
import com.codahale.metrics.MetricRegistry;
import static com.codahale.metrics.MetricRegistry.name;
import com.codahale.metrics.Timer;
import com.codahale.metrics.Timer.Context;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.time.Instant;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;

import static ai.grakn.engine.TaskStatus.FAILED;
import static ai.grakn.engine.TaskStatus.RUNNING;
import static ai.grakn.engine.TaskStatus.STOPPED;
import static ai.grakn.engine.util.ExceptionWrapper.noThrow;
import static java.time.Duration.between;
import static java.time.Instant.now;

/**
 * The {@link SingleQueueTaskRunner} is used by the {@link SingleQueueTaskManager} to execute tasks from a Kafka queue.
 *
 * @author aelred, alexandrorth
 */
public class SingleQueueTaskRunner implements Runnable, AutoCloseable {

    private final static Logger LOG = LoggerFactory.getLogger(SingleQueueTaskRunner.class);

    private final Consumer<TaskState, TaskConfiguration> consumer;
    private final SingleQueueTaskManager manager;
    private final TaskStateStorage storage;
    private final ExternalOffsetStorage offsetStorage;

    private final AtomicBoolean wakeUp = new AtomicBoolean(false);
    private final CountDownLatch countDownLatch = new CountDownLatch(1);
    private final EngineID engineID;
    private final GraknEngineConfig engineConfig;
    private final RedisConnection redis;
    private final Timer readRecordsTimer;
    private final Meter stopMeter;
    private final Meter resubmitMeter;
    private MetricRegistry metricRegistry;
    private final Timer executeTimer;

    private TaskId runningTaskId = null;
    private BackgroundTask runningTask = null;

    /**
     * Create a {@link SingleQueueTaskRunner} which retrieves tasks from the given {@param consumer} and uses the given
     * {@param storage} to store and retrieve information about tasks.
     *  @param manager a place to control the lifecycle of tasks
     * @param engineID identifier of the engine this task runner is on
     * @param offsetStorage a place to externally store kafka offsets
     * @param metricRegistry global metric registry
     */
    public SingleQueueTaskRunner(
            SingleQueueTaskManager manager, EngineID engineID, GraknEngineConfig config,
            RedisConnection redis,
            ExternalOffsetStorage offsetStorage,
            Consumer<TaskState, TaskConfiguration> consumer,
            MetricRegistry metricRegistry){
        this.manager = manager;
        this.storage = manager.storage();
        this.consumer = consumer;
        this.engineID = engineID;
        this.engineConfig = config;
        this.redis = redis;
        this.offsetStorage = offsetStorage;
        this.readRecordsTimer = metricRegistry
                .timer(name(SingleQueueTaskRunner.class, "read-records"));
        this.executeTimer = metricRegistry
                .timer(name(SingleQueueTaskRunner.class, "execute"));
        this.stopMeter = metricRegistry
                .meter(name(SingleQueueTaskRunner.class, "stop"));
        this.resubmitMeter = metricRegistry
                .meter(name(SingleQueueTaskRunner.class, "resubmit"));
        this.metricRegistry = metricRegistry;
    }

    /**
     * Poll Kafka for any new tasks. Will not return until {@link SingleQueueTaskRunner#close()} is called.
     * After receiving tasks, accept as many as possible, up to the maximum allowed number of tasks.
     * For each task, follow the workflow based on its type:
     *  - If not created or not in storage:
     *    - Record that this engine is running this task
     *      Record that this task is running
     *    - Send to thread pool for execution:
     *       - Use reflection to retrieve task
     *       - Start from checkpoint if necessary, or from beginning (TODO)
     *       - Record that this engine is no longer running this task
     *         Mark as completed or failed
     *  - Acknowledge message in queue
     */
    @Override
    public void run() {
        LOG.debug("started");
        while (!wakeUp.get()) {
            try {
                // Reading from both the regular consumer and recurring consumer every time means that we will handle
                // recurring tasks regularly, even if there are lots of non-recurring tasks to process.
                readRecords(consumer);
            } catch (Throwable throwable){
                LOG.error("error thrown", throwable);
                assert false; // This should be unreachable, but in production we still handle it for robustness
            }
        }

        countDownLatch.countDown();
        LOG.debug("stopped");
    }

    /**
     * Close connection to Kafka and thread pool.
     *
     * Inform {@link SingleQueueTaskRunner#run()} method to stop and block until it returns.
     */
    @Override
    public void close() throws Exception {
        wakeUp.set(true);
        noThrow(countDownLatch::await, "Error waiting for the TaskRunner loop to finish");
        noThrow(consumer::close, "Error closing the task runner");
    }

    /**
     * Stop the task if it is executing on this machine
     * @param taskId Identifier of the task to stop
     * @return True if the task is stopped
     */
    public boolean stopTask(TaskId taskId) {
        return taskId.equals(runningTaskId) && runningTask.stop();
    }

    /**
     * Read and handle some records from the given consumer
     */
    private void readRecords(Consumer<TaskState, TaskConfiguration> theConsumer) {
        // This TaskRunner should only ever receive one record from each consumer
        Context context = readRecordsTimer.time();
        try{
            ConsumerRecords<TaskState, TaskConfiguration> records = theConsumer.poll(1000);
            for (ConsumerRecord<TaskState, TaskConfiguration> record : records) {
                TaskState task = record.key();
                TaskConfiguration configuration = record.value();
                handleTask(task, configuration);
                offsetStorage.saveOffset(theConsumer, new TopicPartition(record.topic(), record.partition()));
                LOG.trace("{} acknowledged", task.getId());
            }
        } finally {
            context.stop();
        }
    }

    /**
     * Returns whether the task was succesfully handled, or was just re-submitted.
     */
    private boolean handleTask(TaskState taskFromkafka, TaskConfiguration configuration) {
        LOG.debug("{}\treceived", taskFromkafka);

        TaskState latestState = getLatestState(taskFromkafka);

        if (shouldStopTask(latestState)) {
            stopTask(latestState);<commit_msg>Added metric on record size<commit_after>/*
 * Grakn - A Distributed Semantic Database
 * Copyright (C) 2016  Grakn Labs Limited
 *
 * Grakn is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * Grakn is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with Grakn. If not, see <http://www.gnu.org/licenses/gpl.txt>.
 *
 */

package ai.grakn.engine.tasks.manager.singlequeue;

import ai.grakn.engine.GraknEngineConfig;
import ai.grakn.engine.TaskId;
import static ai.grakn.engine.TaskStatus.FAILED;
import static ai.grakn.engine.TaskStatus.RUNNING;
import static ai.grakn.engine.TaskStatus.STOPPED;
import ai.grakn.engine.tasks.BackgroundTask;
import ai.grakn.engine.tasks.ExternalOffsetStorage;
import ai.grakn.engine.tasks.TaskCheckpoint;
import ai.grakn.engine.tasks.TaskConfiguration;
import ai.grakn.engine.tasks.TaskState;
import ai.grakn.engine.tasks.TaskStateStorage;
import ai.grakn.engine.tasks.connection.RedisConnection;
import ai.grakn.engine.util.EngineID;
import static ai.grakn.engine.util.ExceptionWrapper.noThrow;
import com.codahale.metrics.Histogram;
import com.codahale.metrics.Meter;
import com.codahale.metrics.MetricRegistry;
import static com.codahale.metrics.MetricRegistry.name;
import com.codahale.metrics.Timer;
import com.codahale.metrics.Timer.Context;
import static java.time.Instant.now;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicBoolean;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.common.TopicPartition;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * The {@link SingleQueueTaskRunner} is used by the {@link SingleQueueTaskManager} to execute tasks from a Kafka queue.
 *
 * @author aelred, alexandrorth
 */
public class SingleQueueTaskRunner implements Runnable, AutoCloseable {

    private final static Logger LOG = LoggerFactory.getLogger(SingleQueueTaskRunner.class);

    private final Consumer<TaskState, TaskConfiguration> consumer;
    private final SingleQueueTaskManager manager;
    private final TaskStateStorage storage;
    private final ExternalOffsetStorage offsetStorage;

    private final AtomicBoolean wakeUp = new AtomicBoolean(false);
    private final CountDownLatch countDownLatch = new CountDownLatch(1);
    private final EngineID engineID;
    private final GraknEngineConfig engineConfig;
    private final RedisConnection redis;
    private final Timer readRecordsTimer;
    private final Meter stopMeter;
    private final Meter resubmitMeter;
    private final Histogram recordsSize;
    private MetricRegistry metricRegistry;
    private final Timer executeTimer;

    private TaskId runningTaskId = null;
    private BackgroundTask runningTask = null;

    /**
     * Create a {@link SingleQueueTaskRunner} which retrieves tasks from the given {@param consumer} and uses the given
     * {@param storage} to store and retrieve information about tasks.
     *  @param manager a place to control the lifecycle of tasks
     * @param engineID identifier of the engine this task runner is on
     * @param offsetStorage a place to externally store kafka offsets
     * @param metricRegistry global metric registry
     */
    public SingleQueueTaskRunner(
            SingleQueueTaskManager manager, EngineID engineID, GraknEngineConfig config,
            RedisConnection redis,
            ExternalOffsetStorage offsetStorage,
            Consumer<TaskState, TaskConfiguration> consumer,
            MetricRegistry metricRegistry){
        this.manager = manager;
        this.storage = manager.storage();
        this.consumer = consumer;
        this.engineID = engineID;
        this.engineConfig = config;
        this.redis = redis;
        this.offsetStorage = offsetStorage;
        this.readRecordsTimer = metricRegistry
                .timer(name(SingleQueueTaskRunner.class, "read-records"));
        this.executeTimer = metricRegistry
                .timer(name(SingleQueueTaskRunner.class, "execute"));
        this.stopMeter = metricRegistry
                .meter(name(SingleQueueTaskRunner.class, "stop"));
        this.resubmitMeter = metricRegistry
                .meter(name(SingleQueueTaskRunner.class, "resubmit"));
        this.recordsSize = metricRegistry
                .histogram(name(SingleQueueTaskRunner.class, "records-size"));
        this.metricRegistry = metricRegistry;
    }

    /**
     * Poll Kafka for any new tasks. Will not return until {@link SingleQueueTaskRunner#close()} is called.
     * After receiving tasks, accept as many as possible, up to the maximum allowed number of tasks.
     * For each task, follow the workflow based on its type:
     *  - If not created or not in storage:
     *    - Record that this engine is running this task
     *      Record that this task is running
     *    - Send to thread pool for execution:
     *       - Use reflection to retrieve task
     *       - Start from checkpoint if necessary, or from beginning (TODO)
     *       - Record that this engine is no longer running this task
     *         Mark as completed or failed
     *  - Acknowledge message in queue
     */
    @Override
    public void run() {
        LOG.debug("started");
        while (!wakeUp.get()) {
            try {
                // Reading from both the regular consumer and recurring consumer every time means that we will handle
                // recurring tasks regularly, even if there are lots of non-recurring tasks to process.
                readRecords(consumer);
            } catch (Throwable throwable){
                LOG.error("error thrown", throwable);
                assert false; // This should be unreachable, but in production we still handle it for robustness
            }
        }

        countDownLatch.countDown();
        LOG.debug("stopped");
    }

    /**
     * Close connection to Kafka and thread pool.
     *
     * Inform {@link SingleQueueTaskRunner#run()} method to stop and block until it returns.
     */
    @Override
    public void close() throws Exception {
        wakeUp.set(true);
        noThrow(countDownLatch::await, "Error waiting for the TaskRunner loop to finish");
        noThrow(consumer::close, "Error closing the task runner");
    }

    /**
     * Stop the task if it is executing on this machine
     * @param taskId Identifier of the task to stop
     * @return True if the task is stopped
     */
    public boolean stopTask(TaskId taskId) {
        return taskId.equals(runningTaskId) && runningTask.stop();
    }

    /**
     * Read and handle some records from the given consumer
     */
    private void readRecords(Consumer<TaskState, TaskConfiguration> theConsumer) {
        // This TaskRunner should only ever receive one record from each consumer
        Context context = readRecordsTimer.time();
        try{
            ConsumerRecords<TaskState, TaskConfiguration> records = theConsumer.poll(1000);
            int size = 0;
            for (ConsumerRecord<TaskState, TaskConfiguration> record : records) {
                TaskState task = record.key();
                TaskConfiguration configuration = record.value();
                handleTask(task, configuration);
                offsetStorage.saveOffset(theConsumer, new TopicPartition(record.topic(), record.partition()));
                LOG.trace("{} acknowledged", task.getId());
                size++;
            }
            recordsSize.update(size);
        } finally {
            context.stop();
        }
    }

    /**
     * Returns whether the task was succesfully handled, or was just re-submitted.
     */
    private boolean handleTask(TaskState taskFromkafka, TaskConfiguration configuration) {
        LOG.debug("{}\treceived", taskFromkafka);

        TaskState latestState = getLatestState(taskFromkafka);

        if (shouldStopTask(latestState)) {
            stopTask(latestState);
{'fraction_non_alphanumeric': 0.07862939585211903, 'fraction_numerical': 0.0015028554253080854, 'mean_word_length': 3.0164171897633993, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5532769', 'n_tokens_mistral': 4361, 'n_tokens_neox': 4197, 'n_words': 1409}