<commit_before>/*
 * Copyright 2004-present Facebook. All Rights Reserved.
 */
package com.facebook.presto.server;

import com.facebook.presto.OutputBuffers;
import com.facebook.presto.TaskSource;
import com.facebook.presto.execution.SqlTaskManagerStats;
import com.facebook.presto.execution.TaskId;
import com.facebook.presto.execution.BufferResult;
import com.facebook.presto.execution.TaskInfo;
import com.facebook.presto.execution.TaskManager;
import com.facebook.presto.execution.TaskOutput;
import com.facebook.presto.execution.TaskState;
import com.facebook.presto.operator.Page;
import com.facebook.presto.sql.analyzer.Session;
import com.facebook.presto.sql.planner.PlanFragment;
import com.facebook.presto.util.Threads;
import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.Iterables;
import io.airlift.http.server.HttpServerInfo;
import io.airlift.units.DataSize;
import io.airlift.units.DataSize.Unit;
import io.airlift.units.Duration;

import javax.inject.Inject;
import java.net.URI;
import java.util.Collections;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;

import static com.facebook.presto.block.BlockAssertions.createStringsBlock;
import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom;

public class MockTaskManager
        implements TaskManager
{
    private final Executor executor = Executors.newCachedThreadPool(Threads.daemonThreadsNamed("test-%d"));

    private final HttpServerInfo httpServerInfo;
    private final DataSize maxBufferSize;
    private final int initialPages;

    private final ConcurrentMap<TaskId, TaskOutput> tasks = new ConcurrentHashMap<>();

    @Inject
    public MockTaskManager(HttpServerInfo httpServerInfo)
    {
        this(httpServerInfo, new DataSize(100, Unit.MEGABYTE), 12);
    }

    public MockTaskManager(HttpServerInfo httpServerInfo, DataSize maxBufferSize, int initialPages)
    {
        Preconditions.checkNotNull(httpServerInfo, "httpServerInfo is null");
        Preconditions.checkArgument(maxBufferSize.toBytes() > 0, "pageBufferMax must be at least 1");
        Preconditions.checkArgument(initialPages >= 0, "initialPages is negative");
        Preconditions.checkArgument(initialPages <= maxBufferSize.toBytes(), "initialPages is greater than maxBufferSize");
        this.httpServerInfo = httpServerInfo;
        this.maxBufferSize = maxBufferSize;
        this.initialPages = initialPages;
    }

    @Override
    public synchronized List<TaskInfo> getAllTaskInfo(boolean full)
    {
        ImmutableList.Builder<TaskInfo> builder = ImmutableList.builder();
        for (TaskOutput taskOutput : tasks.values()) {
            builder.add(taskOutput.getTaskInfo(false));
        }
        return builder.build();
    }

    @Override
    public void waitForStateChange(TaskId taskId, TaskState currentState, Duration maxWait)
            throws InterruptedException
    {
    }

    @Override
    public synchronized TaskInfo getTaskInfo(TaskId taskId, boolean full)
    {
        Preconditions.checkNotNull(taskId, "taskId is null");

        TaskOutput taskOutput = tasks.get(taskId);
        if (taskOutput == null) {
            throw new NoSuchElementException();
        }
        return taskOutput.getTaskInfo(false);
    }

    @Override
    public synchronized TaskInfo updateTask(Session session, TaskId taskId, PlanFragment ignored, List<TaskSource> sources, OutputBuffers outputIds)
    {
        Preconditions.checkNotNull(session, "session is null");
        Preconditions.checkNotNull(taskId, "taskId is null");
        Preconditions.checkNotNull(sources, "sources is null");
        Preconditions.checkNotNull(outputIds, "outputIds is null");

        TaskOutput taskOutput = tasks.get(taskId);
        if (taskOutput == null) {
            URI location = uriBuilderFrom(httpServerInfo.getHttpUri()).appendPath("v1/task").appendPath(taskId.toString()).build();
            taskOutput = new TaskOutput(taskId, location, maxBufferSize, executor, new SqlTaskManagerStats());
            tasks.put(taskId, taskOutput);

            List<String> data = ImmutableList.of("apple", "banana", "cherry", "date");

            // load initial pages
            for (int i = 0; i < initialPages; i++) {
                try {
                    taskOutput.addPage(new Page(createStringsBlock(Iterables.concat(Collections.nCopies(i + 1, data)))));
                }
                catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    throw Throwables.propagate(e);
                }
            }
            taskOutput.finish();
        }

        for (String outputId : outputIds.getBufferIds()) {
            taskOutput.addResultQueue(outputId);
        }
        if (outputIds.isNoMoreBufferIds()) {
            taskOutput.noMoreResultQueues();
        }

        return taskOutput.getTaskInfo(false);
    }

    @Override
    public BufferResult getTaskResults(TaskId taskId, String outputId, long startingSequenceId, DataSize maxSize, Duration maxWaitTime)
            throws InterruptedException
    {
        Preconditions.checkNotNull(taskId, "taskId is null");
        Preconditions.checkNotNull(outputId, "outputId is null");

        TaskOutput taskOutput;
        synchronized (this) {
            taskOutput = tasks.get(taskId);
            if (taskOutput == null) {
                throw new NoSuchElementException();
            }
        }
        return taskOutput.getResults(outputId, startingSequenceId, maxSize, maxWaitTime);
    }

    @Override
    public synchronized TaskInfo abortTaskResults(TaskId taskId, String outputId)
    {
        Preconditions.checkNotNull(taskId, "taskId is null");
        Preconditions.checkNotNull(outputId, "outputId is null");

        TaskOutput taskOutput = tasks.get(taskId);
        if (taskOutput == null) {
            throw new NoSuchElementException();
        }
        taskOutput.abortResults(outputId);
        return taskOutput.getTaskInfo(false);
    }

    @Override
    public synchronized TaskInfo cancelTask(TaskId taskId)
    {
        Preconditions.checkNotNull(taskId, "taskId is null");

        TaskOutput taskOutput = tasks.get(taskId);
        if (taskOutput != null) {
            taskOutput.cancel();
            return taskOutput.getTaskInfo(false);
        } else {
            return null;
        }
    }
}<commit_msg>Convert MockTaskManager to use new operator framework<commit_after>/*
 * Copyright 2004-present Facebook. All Rights Reserved.
 */
package com.facebook.presto.server;

import com.facebook.presto.OutputBuffers;
import com.facebook.presto.TaskSource;
import com.facebook.presto.client.FailureInfo;
import com.facebook.presto.execution.BufferResult;
import com.facebook.presto.execution.SharedBuffer;
import com.facebook.presto.execution.TaskId;
import com.facebook.presto.execution.TaskInfo;
import com.facebook.presto.execution.TaskManager;
import com.facebook.presto.execution.TaskState;
import com.facebook.presto.execution.TaskStateMachine;
import com.facebook.presto.noperator.TaskContext;
import com.facebook.presto.operator.Page;
import com.facebook.presto.sql.analyzer.Session;
import com.facebook.presto.sql.planner.PlanFragment;
import com.facebook.presto.sql.planner.plan.PlanNodeId;
import com.facebook.presto.util.Threads;
import com.google.common.base.Preconditions;
import com.google.common.base.Throwables;
import com.google.common.collect.ImmutableList;
import com.google.common.collect.ImmutableSet;
import com.google.common.collect.Iterables;
import io.airlift.http.server.HttpServerInfo;
import io.airlift.units.DataSize;
import io.airlift.units.DataSize.Unit;
import io.airlift.units.Duration;
import org.joda.time.DateTime;

import javax.inject.Inject;

import java.net.URI;
import java.util.Collections;
import java.util.List;
import java.util.NoSuchElementException;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
import java.util.concurrent.Executor;
import java.util.concurrent.Executors;
import java.util.concurrent.atomic.AtomicLong;

import static com.facebook.presto.block.BlockAssertions.createStringsBlock;
import static com.facebook.presto.util.Failures.toFailures;
import static com.google.common.base.Preconditions.checkNotNull;
import static io.airlift.http.client.HttpUriBuilder.uriBuilderFrom;
import static io.airlift.units.DataSize.Unit.MEGABYTE;

public class MockTaskManager
        implements TaskManager
{
    private final Executor executor = Executors.newCachedThreadPool(Threads.daemonThreadsNamed("test-%d"));

    private final HttpServerInfo httpServerInfo;
    private final DataSize maxBufferSize;
    private final int initialPages;

    private final ConcurrentMap<TaskId, MockTask> tasks = new ConcurrentHashMap<>();


    @Inject
    public MockTaskManager(HttpServerInfo httpServerInfo)
    {
        this(httpServerInfo, new DataSize(100, Unit.MEGABYTE), 12);
    }

    public MockTaskManager(HttpServerInfo httpServerInfo, DataSize maxBufferSize, int initialPages)
    {
        checkNotNull(httpServerInfo, "httpServerInfo is null");
        Preconditions.checkArgument(maxBufferSize.toBytes() > 0, "pageBufferMax must be at least 1");
        Preconditions.checkArgument(initialPages >= 0, "initialPages is negative");
        Preconditions.checkArgument(initialPages <= maxBufferSize.toBytes(), "initialPages is greater than maxBufferSize");
        this.httpServerInfo = httpServerInfo;
        this.maxBufferSize = maxBufferSize;
        this.initialPages = initialPages;
    }

    @Override
    public synchronized List<TaskInfo> getAllTaskInfo(boolean full)
    {
        ImmutableList.Builder<TaskInfo> builder = ImmutableList.builder();
        for (MockTask task : tasks.values()) {
            builder.add(task.getTaskInfo());
        }
        return builder.build();
    }

    @Override
    public void waitForStateChange(TaskId taskId, TaskState currentState, Duration maxWait)
            throws InterruptedException
    {
    }

    @Override
    public synchronized TaskInfo getTaskInfo(TaskId taskId, boolean full)
    {
        checkNotNull(taskId, "taskId is null");

        MockTask task = tasks.get(taskId);
        if (task == null) {
            throw new NoSuchElementException();
        }
        return task.getTaskInfo();
    }

    @Override
    public synchronized TaskInfo updateTask(Session session, TaskId taskId, PlanFragment ignored, List<TaskSource> sources, OutputBuffers outputBuffers)
    {
        checkNotNull(session, "session is null");
        checkNotNull(taskId, "taskId is null");
        checkNotNull(sources, "sources is null");
        checkNotNull(outputBuffers, "outputBuffers is null");

        MockTask task = tasks.get(taskId);
        if (task == null) {
            task = new MockTask(session, taskId,
                    uriBuilderFrom(httpServerInfo.getHttpUri()).appendPath("v1/task").appendPath(taskId.toString()).build(), maxBufferSize, initialPages, executor
            );
            tasks.put(taskId, task);
        }
        task.addOutputBuffers(outputBuffers);

        return task.getTaskInfo();
    }

    @Override
    public BufferResult getTaskResults(TaskId taskId, String outputId, long startingSequenceId, DataSize maxSize, Duration maxWaitTime)
            throws InterruptedException
    {
        checkNotNull(taskId, "taskId is null");
        checkNotNull(outputId, "outputId is null");

        MockTask task;
        synchronized (this) {
            task = tasks.get(taskId);
        }

        if (task == null) {
            throw new NoSuchElementException();
        }
        return task.getResults(outputId, startingSequenceId, maxSize, maxWaitTime);
    }

    @Override
    public synchronized TaskInfo abortTaskResults(TaskId taskId, String outputId)
    {
        checkNotNull(taskId, "taskId is null");
        checkNotNull(outputId, "outputId is null");

        MockTask task = tasks.get(taskId);
        if (task == null) {
            throw new NoSuchElementException();
        }
        task.abortResults(outputId);
        return task.getTaskInfo();
    }

    @Override
    public synchronized TaskInfo cancelTask(TaskId taskId)
    {
        checkNotNull(taskId, "taskId is null");

        MockTask task = tasks.get(taskId);
        if (task == null) {
            return null;
        }

        task.cancel();
        return task.getTaskInfo();
    }

    public static class MockTask
    {
        private final AtomicLong nextTaskInfoVersion = new AtomicLong(TaskInfo.STARTING_VERSION);

        private final URI location;
        private final TaskStateMachine taskStateMachine;
        private final TaskContext taskContext;
        private final SharedBuffer sharedBuffer;

        public MockTask(Session session,
                TaskId taskId,
                URI location,
                DataSize maxBufferSize,
                int initialPages,
                Executor executor)
        {
            this.taskStateMachine = new TaskStateMachine(checkNotNull(taskId, "taskId is null"), checkNotNull(executor, "executor is null"));
            this.taskContext = new TaskContext(taskStateMachine, executor, session, new DataSize(256, MEGABYTE), new DataSize(1, MEGABYTE));

            this.location = checkNotNull(location, "location is null");

            this.sharedBuffer = new SharedBuffer(checkNotNull(maxBufferSize, "maxBufferSize is null"));

            List<String> data = ImmutableList.of("apple", "banana", "cherry", "date");

            // load initial pages
            for (int i = 0; i < initialPages; i++) {
                try {
                    sharedBuffer.add(new Page(createStringsBlock(Iterables.concat(Collections.nCopies(i + 1, data)))));
                }
                catch (InterruptedException e) {
                    Thread.currentThread().interrupt();
                    throw Throwables.propagate(e);
                }
            }
            sharedBuffer.finish();
        }

        public void abortResults(String outputId)
        {
            sharedBuffer.abort(outputId);
        }

        public void addOutputBuffers(OutputBuffers outputBuffers)
        {
            for (String outputId : outputBuffers.getBufferIds()) {
                sharedBuffer.addQueue(outputId);
            }
            if (outputBuffers.isNoMoreBufferIds()) {
                sharedBuffer.noMoreQueues();
            }
        }

        public void cancel()
        {
            taskStateMachine.cancel();
        }

        public BufferResult getResults(String outputId, long startingSequenceId, DataSize maxSize, Duration maxWaitTime)
                throws InterruptedException
        {
            return sharedBuffer.get(outputId, startingSequenceId, maxSize, maxWaitTime);
        }

        public TaskInfo getTaskInfo()
        {
            TaskState state = taskStateMachine.getState();
            List<FailureInfo> failures = ImmutableList.of();
            if (state == TaskState.FAILED) {
                failures = toFailures(taskStateMachine.getFailureCauses());
            }

            return new TaskInfo(
                    taskStateMachine.getTaskId(),
                    nextTaskInfoVersion.getAndIncrement(),
                    state,
                    location,
                    DateTime.now(),
                    sharedBuffer.getInfo(),
                    ImmutableSet.<PlanNodeId>of(),
                    taskContext.getTaskStats(),
                    failures,
                    taskContext.getOutputItems());
        }
    }
}
{'fraction_non_alphanumeric': 0.0925681203196778, 'fraction_numerical': 0.002139575860549997, 'mean_word_length': 2.8657261007054244, 'pattern_counts': {'":': 0, '<': 21, '<?xml version=': 0, '>': 21, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4942512', 'n_tokens_mistral': 4307, 'n_tokens_neox': 4208, 'n_words': 906}