All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.trino.execution.executor.scheduler.TaskControl Maven / Gradle / Ivy

/*
 * 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.execution.executor.scheduler;

import com.google.common.base.Ticker;
import com.google.errorprone.annotations.ThreadSafe;
import com.google.errorprone.annotations.concurrent.GuardedBy;

import java.util.Objects;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import static java.util.Objects.requireNonNull;

/**
 * Equality is based on group and id for the purpose of adding to the scheduling queue.
 */
@ThreadSafe
final class TaskControl
{
    private final Group group;
    private final int id;
    private final Ticker ticker;

    private final Lock lock = new ReentrantLock();

    @GuardedBy("lock")
    private final Condition wakeup = lock.newCondition();

    @GuardedBy("lock")
    private boolean ready;

    @GuardedBy("lock")
    private boolean blocked;

    @GuardedBy("lock")
    private boolean cancelled;

    @GuardedBy("lock")
    private State state;

    private volatile long periodStart;
    private final AtomicLong startNanos = new AtomicLong();
    private final AtomicLong scheduledNanos = new AtomicLong();
    private final AtomicLong blockedNanos = new AtomicLong();
    private final AtomicLong waitNanos = new AtomicLong();
    private volatile Thread thread;

    public TaskControl(Group group, int id, Ticker ticker)
    {
        this.group = requireNonNull(group, "group is null");
        this.id = id;
        this.ticker = requireNonNull(ticker, "ticker is null");
        this.state = State.NEW;
        this.ready = false;
        this.periodStart = ticker.read();
    }

    public int id()
    {
        return id;
    }

    public void setThread(Thread thread)
    {
        this.thread = thread;
    }

    public void cancel()
    {
        lock.lock();
        try {
            cancelled = true;
            wakeup.signal();

            // TODO: it should be possible to interrupt the thread, but
            //       it appears that it's not safe to do so. It can cause the query
            //       to get stuck (e.g., AbstractDistributedEngineOnlyQueries.testSelectiveLimit)
            //
            //       Thread thread = this.thread;
            //       if (thread != null) {
            //           thread.interrupt();
            //       }
        }
        finally {
            lock.unlock();
        }
    }

    /**
     * Called by the scheduler thread when the task is ready to run. It
     * causes anyone blocking in {@link #awaitReady()} to wake up.
     *
     * @return false if the task was already cancelled
     */
    public boolean markReady()
    {
        lock.lock();
        try {
            if (cancelled) {
                return false;
            }
            ready = true;
            wakeup.signal();
        }
        finally {
            lock.unlock();
        }

        return true;
    }

    public void markNotReady()
    {
        lock.lock();
        try {
            ready = false;
        }
        finally {
            lock.unlock();
        }
    }

    public boolean isReady()
    {
        lock.lock();
        try {
            return ready;
        }
        finally {
            lock.unlock();
        }
    }

    /**
     * @return false if the operation was interrupted due to cancellation
     */
    public boolean awaitReady()
    {
        lock.lock();
        try {
            while (!ready && !cancelled) {
                try {
                    wakeup.await();
                }
                catch (InterruptedException e) {
                }
            }

            return !cancelled;
        }
        finally {
            lock.unlock();
        }
    }

    public void markUnblocked()
    {
        lock.lock();
        try {
            blocked = false;
            wakeup.signal();
        }
        finally {
            lock.unlock();
        }
    }

    public void markBlocked()
    {
        lock.lock();
        try {
            blocked = true;
        }
        finally {
            lock.unlock();
        }
    }

    public void awaitUnblock()
    {
        lock.lock();
        try {
            while (blocked && !cancelled) {
                try {
                    wakeup.await();
                }
                catch (InterruptedException e) {
                }
            }
        }
        finally {
            lock.unlock();
        }
    }

    /**
     * @return false if the transition was unsuccessful due to the task being interrupted
     */
    public boolean transitionToBlocked()
    {
        boolean success = transitionTo(State.BLOCKED);

        if (success) {
            markBlocked();
        }

        return success;
    }

    public void transitionToFinished()
    {
        transitionTo(State.FINISHED);
    }

    /**
     * @return false if the transition was unsuccessful due to the task being interrupted
     */
    public boolean transitionToWaiting()
    {
        boolean success = transitionTo(State.WAITING);

        if (success) {
            markNotReady();
        }

        return success;
    }

    /**
     * @return false if the transition was unsuccessful due to the task being interrupted
     */
    public boolean transitionToRunning()
    {
        return transitionTo(State.RUNNING);
    }

    private boolean transitionTo(State state)
    {
        lock.lock();
        try {
            recordPeriodEnd(this.state);

            if (cancelled) {
                this.state = State.INTERRUPTED;
                return false;
            }
            else {
                this.state = state;
                return true;
            }
        }
        finally {
            lock.unlock();
        }
    }

    private void recordPeriodEnd(State state)
    {
        long now = ticker.read();
        long elapsed = now - periodStart;
        switch (state) {
            case RUNNING -> scheduledNanos.addAndGet(elapsed);
            case BLOCKED -> blockedNanos.addAndGet(elapsed);
            case NEW -> startNanos.addAndGet(elapsed);
            case WAITING -> waitNanos.addAndGet(elapsed);
            case INTERRUPTED, FINISHED -> {}
        }
        periodStart = now;
    }

    public Group group()
    {
        return group;
    }

    public State getState()
    {
        lock.lock();
        try {
            return state;
        }
        finally {
            lock.unlock();
        }
    }

    public long elapsed()
    {
        return ticker.read() - periodStart;
    }

    public long getStartNanos()
    {
        return startNanos.get();
    }

    public long getWaitNanos()
    {
        return waitNanos.get();
    }

    public long getScheduledNanos()
    {
        return scheduledNanos.get();
    }

    public long getBlockedNanos()
    {
        return blockedNanos.get();
    }

    @Override
    public boolean equals(Object o)
    {
        if (this == o) {
            return true;
        }
        if (o == null || getClass() != o.getClass()) {
            return false;
        }
        TaskControl that = (TaskControl) o;
        return id == that.id && group.equals(that.group);
    }

    @Override
    public int hashCode()
    {
        return Objects.hash(group, id);
    }

    @Override
    public String toString()
    {
        lock.lock();
        try {
            return "%s-%s [%s, %sready, %sblocked, %scancelled, start=%s, tid=%s]".formatted(
                    group.name(),
                    id,
                    state,
                    ready ? "+" : "-",
                    blocked ? "+" : "-",
                    cancelled ? "+" : "-",
                    periodStart,
                    thread.threadId());
        }
        finally {
            lock.unlock();
        }
    }

    public Thread getThread()
    {
        return thread;
    }

    public enum State
    {
        NEW,
        WAITING,
        RUNNING,
        BLOCKED,
        INTERRUPTED,
        FINISHED
    }
}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy