ai.grakn.engine.tasks.manager.redisqueue.RedisTaskManager Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of grakn-engine Show documentation
There is a newer version: 1.4.3
/*
 * 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 .
 *
 */

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

import ai.grakn.GraknConfigKey;
import ai.grakn.engine.GraknConfig;
import ai.grakn.engine.TaskId;
import ai.grakn.engine.factory.EngineGraknTxFactory;
import ai.grakn.engine.postprocessing.PostProcessor;
import ai.grakn.engine.tasks.manager.TaskConfiguration;
import ai.grakn.engine.tasks.manager.TaskManager;
import ai.grakn.engine.tasks.manager.TaskState;
import ai.grakn.engine.util.EngineID;
import ai.grakn.redisq.Redisq;
import ai.grakn.redisq.RedisqBuilder;
import ai.grakn.redisq.exceptions.StateFutureInitializationException;
import ai.grakn.redisq.exceptions.WaitException;
import com.codahale.metrics.MetricRegistry;
import com.google.common.collect.ImmutableSet;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import redis.clients.jedis.Jedis;
import redis.clients.util.Pool;

import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.function.Consumer;

import static ai.grakn.redisq.State.DONE;
import static ai.grakn.redisq.State.FAILED;

/**
 * Handle the lifecycle of tasks in Redis. Given a jedis pool
 * it starts a set of consumers that subscribe to the task queue.
 * Tasks can be added using the addTask method.
 *
 * @author pluraliseseverythings
 */
public class RedisTaskManager implements TaskManager {
    private static final Logger LOG = LoggerFactory.getLogger(RedisTaskManager.class);
    private static final String QUEUE_NAME = "grakn";

    private final Redisq redisq;
    private final RedisTaskStorage taskStorage;

    public RedisTaskManager(EngineID engineId, GraknConfig config, Pool jedisPool,
                            int threads, EngineGraknTxFactory factory,
                            MetricRegistry metricRegistry, PostProcessor postProcessor) {

        Consumer consumer = new RedisTaskQueueConsumer(this, engineId, config,
                metricRegistry, factory,
                postProcessor);

        LOG.info("Running queue consumer with {} execution threads", threads);
        this.redisq = new RedisqBuilder()
                .setJedisPool(jedisPool)
                .setName(QUEUE_NAME)
                .setConsumer(consumer)
                .setMetricRegistry(metricRegistry)
                .setThreadPoolSize(threads)
                .setDelay(config.getProperty(GraknConfigKey.TASK_DELAY))
                .setDocumentClass(Task.class)
                .createRedisq();
        this.taskStorage = RedisTaskStorage.create(redisq, metricRegistry);
    }

    @Override
    public void close() {
        LOG.info("Closing task manager");
        try {
            redisq.close();
        } catch (InterruptedException e) {
            LOG.error("Interrupted while closing queue", e);
        }
    }

    @Override
    public CompletableFuture start() {
        return CompletableFuture
                .runAsync(redisq::startConsumer)
                .exceptionally(e -> {
                    close();
                    throw new RuntimeException("Failed to initialise subscription");
                });
    }

    @Override
    public RedisTaskStorage storage() {
        return taskStorage;
    }

    @Override
    public void addTask(TaskState taskState, TaskConfiguration configuration) {
        Task task = Task.builder()
                .setTaskConfiguration(configuration)
                .setTaskState(taskState).build();
        redisq.push(task);
    }

    @Override
    public void runTask(TaskState taskState, TaskConfiguration configuration) {
        Task task = Task.builder()
                .setTaskConfiguration(configuration)
                .setTaskState(taskState).build();
        try {
            redisq.pushAndWait(task, 5, TimeUnit.MINUTES);
        } catch (WaitException e) {
            throw new RuntimeException("Could not run task", e);
        }
    }

    public Future subscribeToTask(TaskId taskId)
            throws StateFutureInitializationException, ExecutionException, InterruptedException {
        return redisq.getFutureForDocumentStateWait(ImmutableSet.of(DONE, FAILED), taskId.value());
    }


    public void waitForTask(TaskId taskId, long timeout, TimeUnit timeUnit)
            throws StateFutureInitializationException, ExecutionException, InterruptedException, TimeoutException {
        redisq.getFutureForDocumentStateWait(ImmutableSet.of(DONE, FAILED), taskId.value()).get(timeout, timeUnit);
    }

    public Redisq getQueue() {
        return redisq;
    }
}