/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you 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 org.apache.reef.vortex.api;

import org.apache.reef.annotations.Unstable;
import org.apache.reef.wake.EventHandler;
import org.apache.reef.wake.impl.ThreadPoolStage;

import java.util.concurrent.*;

/**
 * The interface between user code and submitted task.
 */
@Unstable
public final class VortexFuture<TOutput> implements Future<TOutput> {
  private TOutput userResult;
  private Exception userException;
  private final CountDownLatch countDownLatch = new CountDownLatch(1);
  private final ThreadPoolStage<TOutput> stage;

  /**
   * Creates a {@link VortexFuture}.
   */
  public VortexFuture() {
    stage = null;
  }

  /**
   * Creates a {@link VortexFuture} with a callback.
   */
  public VortexFuture(final EventHandler<TOutput> callbackHandler) {
    stage = new ThreadPoolStage<>(callbackHandler, 1);
  }

  /**
   * TODO[REEF-502]: Support Vortex Tasklet(s) cancellation by user.
   */
  @Override
  public boolean cancel(final boolean mayInterruptIfRunning) {
    throw new UnsupportedOperationException("Cancel not yet supported");
  }

  /**
   * TODO[REEF-502]: Support Vortex Tasklet(s) cancellation by user.
   */
  @Override
  public boolean isCancelled() {
    throw new UnsupportedOperationException("Cancel not yet supported");
  }

  /**
   * @return true it the task completed, false if not.
   */
  @Override
  public boolean isDone() {
    return countDownLatch.getCount() == 0;
  }

  /**
   * Infinitely wait for the result of the task.
   */
  @Override
  public TOutput get() throws InterruptedException, ExecutionException {
    countDownLatch.await();
    if (userResult != null) {
      return userResult;
    } else {
      assert userException != null;
      throw new ExecutionException(userException);
    }
  }

  /**
   * Wait a certain period of time for the result of the task.
   */
  @Override
  public TOutput get(final long timeout, final TimeUnit unit)
      throws InterruptedException, ExecutionException, TimeoutException {
    if (!countDownLatch.await(timeout, unit)) {
      throw new TimeoutException();
    }

    if (userResult != null) {
      return userResult;
    } else {
      assert userException != null;
      throw new ExecutionException(userException);
    }
  }

  /**
   * Called by VortexMaster to let the user know that the task completed.
   */
  public void completed(final TOutput result) {
    this.userResult = result;
    if (stage != null) {
      stage.onNext(userResult);
    }
    this.countDownLatch.countDown();
  }

  /**
   * Called by VortexMaster to let the user know that the task threw an exception.
   */
  public void threwException(final Exception exception) {
    this.userException = exception;
    this.countDownLatch.countDown();
  }
}

@@ -883,16 +883,54 @@
 stable;%0A
+import org.apache.reef.util.Optional;%0A
 import o
@@ -1195,23 +1195,164 @@
 %7B%0A  
-private 
+// userResult starts out as null. If not null =%3E variable is set and tasklet returned.%0A  // Otherwise tasklet has not completed.%0A  private Optional%3C
 TOutput
+%3E
  use
@@ -1354,24 +1354,31 @@
 %3E userResult
+ = null
 ;%0A  private 
@@ -2626,32 +2626,38 @@
 eturn userResult
+.get()
 ;%0A    %7D else %7B%0A 
@@ -3116,24 +3116,30 @@
 n userResult
+.get()
 ;%0A    %7D else
@@ -3395,22 +3395,43 @@
 esult = 
+Optional.ofNullable(
 result
+)
 ;%0A    if
@@ -3474,24 +3474,30 @@
 t(userResult
+.get()
 );%0A    %7D%0A   

Use Optional in favor of null check in VortexFuture