com.linkedin.parseq.BaseTask Maven / Gradle / Ivy

Go to download
/*
 * Copyright 2012 LinkedIn, Inc
 *
 * 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 com.linkedin.parseq;

import java.util.Collection;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;
import java.util.function.Function;
import java.util.function.Supplier;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import com.linkedin.parseq.internal.IdGenerator;
import com.linkedin.parseq.internal.TaskLogger;
import com.linkedin.parseq.promise.DelegatingPromise;
import com.linkedin.parseq.promise.Promise;
import com.linkedin.parseq.promise.Promises;
import com.linkedin.parseq.promise.SettablePromise;
import com.linkedin.parseq.trace.Relationship;
import com.linkedin.parseq.trace.ResultType;
import com.linkedin.parseq.trace.ShallowTrace;
import com.linkedin.parseq.trace.ShallowTraceBuilder;
import com.linkedin.parseq.trace.Trace;
import com.linkedin.parseq.trace.TraceBuilder;


/**
 * An abstract base class that can be used to build implementations of
 * {@link Task}.
 *
 * @author Chris Pettitt ([email protected])
 * @author Chi Chan ([email protected])
 */
public abstract class BaseTask extends DelegatingPromiseimplements Task {
  static final Logger LOGGER = LoggerFactory.getLogger(BaseTask.class);

  private static enum StateType {
    // The initial state of the task.
    //
    // A task in this state can be cancelled and have its priority changed.
    INIT,

    // The task is currently executing. That is, a thread is in the run()
    // method for the TaskDef.
    //
    // A task in this state in not cancellable and cannot have its priority
    // changed.
    RUN,

    // The task has finished running, but the result has not yet been set.
    // This occurs for Tasks with AsyncTaskDefs.
    //
    // A task in this state in cancellable, and cannot have its priority
    // changed.
    PENDING,

    // The task is resolved.
    //
    // A task in this state in not cancellable and cannot have its priority
    // changed.
    DONE
  }

  private final Long _id = IdGenerator.getNextId();
  private final AtomicReference _stateRef;
  private final String _name;
  protected final ShallowTraceBuilder _shallowTraceBuilder;

  protected volatile Function _traceValueProvider;

  private volatile TraceBuilder _traceBuilder;

  /**
   * Constructs a base task without a specified name. The name for this task
   * will be the {@link #toString} representation for this instance. It is
   * usually best to use {@link BaseTask#BaseTask(String)}.
   */
  public BaseTask() {
    this(null);
  }

  /**
   * Constructs a base task with a name.
   *
   * @param name the name for this task.
   */
  public BaseTask(final String name) {
    super(Promises. settable());
    _name = name;
    final State state = State.INIT;
    _shallowTraceBuilder = new ShallowTraceBuilder(_id);
    _shallowTraceBuilder.setName(getName());
    _shallowTraceBuilder.setResultType(ResultType.UNFINISHED);
    _stateRef = new AtomicReference(state);
  }

  @Override
  public Long getId() {
    return _id;
  }

  @Override
  public int getPriority() {
    return _stateRef.get().getPriority();
  }

  @Override
  public boolean setPriority(final int priority) {
    if (priority < Priority.MIN_PRIORITY || priority > Priority.MAX_PRIORITY) {
      throw new IllegalArgumentException("Priority out of bounds: " + priority);
    }

    State state;
    State newState;
    do {
      state = _stateRef.get();
      if (state.getType() != StateType.INIT) {
        return false;
      }

      newState = new State(state.getType(), priority);
    } while (!_stateRef.compareAndSet(state, newState));

    return true;
  }

  @Override
  public TraceBuilder getTraceBuilder() {
    return _traceBuilder;
  }

  @Override
  public final void contextRun(final Context context, final Task parent, final Collection> predecessors) {
    final TaskLogger taskLogger = context.getTaskLogger();
    final TraceBuilder traceBuilder = context.getTraceBuilder();
    if (transitionRun(traceBuilder)) {
      final Promise promise;
      try {
        if (parent != null) {
          traceBuilder.addRelationship(Relationship.CHILD_OF, getShallowTraceBuilder(),
              parent.getShallowTraceBuilder());
        }
        for (Task predecessor : predecessors) {
          traceBuilder.addRelationship(Relationship.SUCCESSOR_OF, getShallowTraceBuilder(),
              predecessor.getShallowTraceBuilder());
        }

        taskLogger.logTaskStart(this);
        try {
          final Context wrapperContext = new WrappedContext(context);
          promise = doContextRun(wrapperContext);
        } finally {
          transitionPending();
        }

        promise.addListener(resolvedPromise -> {
          if (resolvedPromise.isFailed()) {
            fail(resolvedPromise.getError(), taskLogger);
          } else {
            done(resolvedPromise.get(), taskLogger);
          }
        } );
      } catch (Throwable t) {
        fail(t, taskLogger);
      }
    } else {
      //this is possible when task was cancelled or has been executed multiple times
      //e.g. task has multiple paths it can be executed with or has been completed by
      //a FusionTask
      if (parent != null) {
        traceBuilder.addRelationship(Relationship.POTENTIAL_CHILD_OF, getShallowTraceBuilder(),
            parent.getShallowTraceBuilder());
      }
      for (Task predecessor : predecessors) {
        traceBuilder.addRelationship(Relationship.POSSIBLE_SUCCESSOR_OF, getShallowTraceBuilder(),
            predecessor.getShallowTraceBuilder());
      }
    }
  }

  @SuppressWarnings("unchecked")
  private Promise doContextRun(final Context context) throws Throwable {
    return (Promise) run(context);
  }

  /**
   * Returns the name of this task. If no name was set during construction
   * this method will return the value of {@link #toString()}. In most
   * cases it is preferable to explicitly set a name.
   *
   * @return the name of this task
   */
  @Override
  public String getName() {
    return _name == null ? toString() : _name;
  }

  @Override
  public boolean cancel(final Exception rootReason) {
    if (transitionCancel(rootReason)) {
      final Exception reason = new CancellationException(rootReason);
      traceFailure(reason);
      getSettableDelegate().fail(reason);
      return true;
    }
    return false;
  }

  protected void traceFailure(final Throwable reason) {
    if (Exceptions.isEarlyFinish(reason)) {
      _shallowTraceBuilder.setResultType(ResultType.EARLY_FINISH);
    } else {
      _shallowTraceBuilder.setResultType(ResultType.ERROR);
      _shallowTraceBuilder.setValue(Exceptions.failureToString(reason));
    }
  }

  @Override
  public ShallowTraceBuilder getShallowTraceBuilder() {
    return _shallowTraceBuilder;
  }

  @Override
  public ShallowTrace getShallowTrace() {
    return _shallowTraceBuilder.build();
  }

  @Override
  public void setTraceValueSerializer(final Function traceValueProvider) {
    _traceValueProvider = traceValueProvider;
  }

  @Override
  public Trace getTrace() {
    TraceBuilder traceBuilder = getTraceBuilder();
    if (traceBuilder != null) {
      return traceBuilder.build();
    } else {
      return Trace.single(getShallowTrace());
    }
  }

  /**
   * This template method is invoked when the task is run.
   *
   * @param context the context to use while running this task
   * @return a promise that will have its value set when this task is finished
   * @throws Exception if an error occurs while running the task
   */
  protected abstract Promise run(final Context context) throws Throwable;

  private void traceDone(final T value) {
    _shallowTraceBuilder.setResultType(ResultType.SUCCESS);
    final Function traceValueProvider = _traceValueProvider;
    if (traceValueProvider != null) {
      try {
        _shallowTraceBuilder.setValue(traceValueProvider.apply(value));
      } catch (Exception e) {
        _shallowTraceBuilder.setValue(Exceptions.failureToString(e));
      }
    }
  }

  private void done(final T value, final TaskLogger taskLogger) {
    if (transitionDone()) {
      traceDone(value);
      getSettableDelegate().done(value);
      taskLogger.logTaskEnd(BaseTask.this, _traceValueProvider);
    }
  }

  private void fail(final Throwable error, final TaskLogger taskLogger) {
    if (transitionDone()) {
      traceFailure(error);
      getSettableDelegate().fail(error);
      taskLogger.logTaskEnd(BaseTask.this, _traceValueProvider);
    }
  }

  protected boolean transitionRun(final TraceBuilder traceBuilder) {
    State state;
    State newState;
    final long startNanos = System.nanoTime();
    do {
      state = _stateRef.get();
      if (state.getType() != StateType.INIT) {
        return false;
      }
      newState = state.transitionRun();
    } while (!_stateRef.compareAndSet(state, newState));
    _traceBuilder = traceBuilder;
    traceBuilder.addShallowTrace(_shallowTraceBuilder);
    _shallowTraceBuilder.setStartNanos(startNanos);

    return true;
  }

  protected void transitionPending() {
    State state;
    State newState;
    final long pendingNanos = System.nanoTime();
    do {
      state = _stateRef.get();
      if (state.getType() != StateType.RUN) {
        return;
      }
      newState = state.transitionPending();
    } while (!_stateRef.compareAndSet(state, newState));
    _shallowTraceBuilder.setPendingNanos(pendingNanos);
  }

  protected boolean transitionCancel(final Exception reason) {
    State state;
    State newState;

    do {
      state = _stateRef.get();
      final StateType type = state.getType();
      if (type == StateType.RUN || type == StateType.DONE) {
        return false;
      }

      newState = state.transitionDone();
    } while (!_stateRef.compareAndSet(state, newState));

    return true;
  }

  protected boolean transitionDone() {
    State state;
    State newState;
    do {
      state = _stateRef.get();
      if (state.getType() == StateType.DONE) {
        return false;
      }

      newState = state.transitionDone();
    } while (!_stateRef.compareAndSet(state, newState));
    return true;
  }

  protected SettablePromise getSettableDelegate() {
    return (SettablePromise) super.getDelegate();
  }

  protected static class State {
    private final StateType _type;
    private final int _priority;

    private State(final StateType type, final int priority) {
      _type = type;
      _priority = priority;
    }

    public StateType getType() {
      return _type;
    }

    public int getPriority() {
      return _priority;
    }

    public State transitionRun() {
      return new State(StateType.RUN, _priority);
    }

    public State transitionPending() {
      return new State(StateType.PENDING, _priority);
    }

    public State transitionDone() {
      return new State(StateType.DONE, _priority);
    }

    public static final State INIT = new State(StateType.INIT, Priority.DEFAULT_PRIORITY) {
      @Override
      final public State transitionDone() {
        return DONE;
      };

      @Override
      final public State transitionRun() {
        return RUN;
      }

      @Override
      final public State transitionPending() {
        return PENDING;
      }
    };

    public static final State RUN = new State(StateType.RUN, Priority.DEFAULT_PRIORITY) {
      @Override
      final public State transitionDone() {
        return DONE;
      };

      @Override
      final public State transitionRun() {
        return RUN;
      }

      @Override
      final public State transitionPending() {
        return PENDING;
      }
    };

    public static final State PENDING = new State(StateType.PENDING, Priority.DEFAULT_PRIORITY) {
      @Override
      final public State transitionDone() {
        return DONE;
      };

      @Override
      final public State transitionRun() {
        return RUN;
      }

      @Override
      final public State transitionPending() {
        return PENDING;
      }
    };

    public static final State DONE = new State(StateType.DONE, Priority.DEFAULT_PRIORITY) {
      @Override
      final public State transitionDone() {
        return DONE;
      };

      @Override
      final public State transitionRun() {
        return RUN;
      }

      @Override
      final public State transitionPending() {
        return PENDING;
      }
    };
  }

  private class WrappedContext implements Context {
    private final Context _context;

    public WrappedContext(final Context context) {
      _context = context;
    }

    @Override
    public Cancellable createTimer(final long time, final TimeUnit unit, final Task task) {
      final Cancellable cancellable = _context.createTimer(time, unit, task);
      getTraceBuilder().addRelationship(Relationship.POTENTIAL_PARENT_OF, getShallowTraceBuilder(),
          task.getShallowTraceBuilder());
      return cancellable;
    }

    @Override
    public void run(final Task... tasks) {
      _context.run(tasks);
      for (Task task : tasks) {
        getTraceBuilder().addRelationship(Relationship.POTENTIAL_PARENT_OF, getShallowTraceBuilder(),
            task.getShallowTraceBuilder());
      }
    }

    @Override
    public After after(final Promise... promises) {
      return new WrappedAfter(_context.after(promises));
    }

    @Override
    public Object getEngineProperty(String key) {
      return _context.getEngineProperty(key);
    }

    @Override
    public TraceBuilder getTraceBuilder() {
      return _context.getTraceBuilder();
    }

    private class WrappedAfter implements After {
      private final After _after;

      public WrappedAfter(final After after) {
        _after = after;
      }

      @Override
      public void run(final Task task) {
        _after.run(task);
        getTraceBuilder().addRelationship(Relationship.POTENTIAL_PARENT_OF, getShallowTraceBuilder(),
            task.getShallowTraceBuilder());
      }

      @Override
      public void run(Supplier> taskSupplier) {
        _after.run(() -> {
          Task task = taskSupplier.get();
          if (task != null) {
            getTraceBuilder().addRelationship(Relationship.POTENTIAL_PARENT_OF, getShallowTraceBuilder(),
                task.getShallowTraceBuilder());
          }
          return task;
        } );
      }

      @Override
      public void runSideEffect(final Task task) {
        _after.runSideEffect(task);
        getTraceBuilder().addRelationship(Relationship.POTENTIAL_PARENT_OF, getShallowTraceBuilder(),
            task.getShallowTraceBuilder());
      }
    }

    @Override
    public ShallowTraceBuilder getShallowTraceBuilder() {
      return _context.getShallowTraceBuilder();
    }

    @Override
    public Long getPlanId() {
      return _context.getPlanId();
    }

    @Override
    public Long getTaskId() {
      return _context.getTaskId();
    }

    @Override
    public TaskLogger getTaskLogger() {
      return _context.getTaskLogger();
    }

    @Override
    public void runSideEffect(Task... tasks) {
      _context.runSideEffect(tasks);
      for (Task task : tasks) {
        getTraceBuilder().addRelationship(Relationship.POTENTIAL_PARENT_OF, getShallowTraceBuilder(),
            task.getShallowTraceBuilder());
      }
    }
  }

  @Override
  public String toString() {
    return "Task [id=" + _id + ", name=" + _name + "]";
  }

}