/*
 * 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.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicReference;

import static org.testng.Assert.*;

import org.testng.annotations.Test;

import com.linkedin.parseq.promise.Promises;


/**
 * @author Jaroslaw Odzga (jodzga@linkedin.com)
 */
public class TestTaskCancellation extends BaseEngineTest {

  @Test
  public void testTaskCancellationAfterRun() throws InterruptedException {
    final AtomicReference<Throwable> cancelActionValue = new AtomicReference<>();
    final CountDownLatch runLatch = new CountDownLatch(1);
    final CountDownLatch listenerLatch = new CountDownLatch(1);
    Task<?> uncompleted = Task.async(() -> {
      runLatch.countDown();
      return Promises.settable();
    });
    uncompleted.addListener(p -> {
      if (p.isFailed() && Exceptions.isCancellation(p.getError())) {
        cancelActionValue.set(p.getError().getCause());
      }
      listenerLatch.countDown();
    } );
    getEngine().run(uncompleted);
    assertTrue(runLatch.await(5, TimeUnit.SECONDS));
    Exception cancelReason = new Exception();
    assertTrue(uncompleted.cancel(cancelReason));
    assertTrue(listenerLatch.await(5, TimeUnit.SECONDS));
    logTracingResults("TestTaskCancellation.testTaskCancellationAfterRun", uncompleted);
    assertEquals(cancelActionValue.get(), cancelReason);
  }

  @Test
  public void testTaskCancellationBeforeRun() throws InterruptedException {
    final AtomicReference<Throwable> cancelActionValue = new AtomicReference<>();
    Task<?> uncompleted = Task.async(() -> Promises.settable());
    uncompleted.addListener(p -> {
      if (p.isFailed() && Exceptions.isCancellation(p.getError())) {
        cancelActionValue.set(p.getError().getCause());
      }
    } );
    Exception cancelReason = new Exception();
    assertTrue(uncompleted.cancel(cancelReason));
    getEngine().run(uncompleted);
    uncompleted.await(5, TimeUnit.SECONDS);
    logTracingResults("TestTaskCancellation.testTaskCancellationBeforeRun", uncompleted);
    assertEquals(cancelActionValue.get(), cancelReason);
  }

  @Test
  public void testTaskCancellationAfterCompleted() throws InterruptedException {
    final AtomicReference<Throwable> cancelActionValue = new AtomicReference<>();
    Task<?> completed = Task.value(10);
    completed.addListener(p -> {
      if (p.isFailed() && Exceptions.isCancellation(p.getError())) {
        cancelActionValue.set(p.getError().getCause());
      }
    } );
    runAndWait("TestTaskCancellation.testTaskCancellationAfterCompleted", completed);
    Exception cancelReason = new Exception();
    assertFalse(completed.cancel(cancelReason));
    assertNull(cancelActionValue.get());
  }

  @Test
  public void testTaskCancellationPar() throws InterruptedException {
    final AtomicReference<Throwable> cancelActionValue = new AtomicReference<>();
    Task<Integer> completed = Task.value(10);
    final CountDownLatch runLatch = new CountDownLatch(1);
    final CountDownLatch listenerLatch = new CountDownLatch(1);
    Task<Integer> uncompleted = Task.async(() -> {
      runLatch.countDown();
      return Promises.settable();
    });
    uncompleted.addListener(p -> {
      if (p.isFailed() && Exceptions.isCancellation(p.getError())) {
        cancelActionValue.set(p.getError().getCause());
      }
      listenerLatch.countDown();
    } );
    Task<?> task =
        Task.par(completed, uncompleted).map((x, y) -> x + y).withTimeout(100, TimeUnit.MILLISECONDS).recover(e -> 0);
    runAndWait("TestTaskCancellation.testTaskCancellationPar", task);
    assertTrue(listenerLatch.await(5, TimeUnit.SECONDS));

    assertTrue(cancelActionValue.get() instanceof EarlyFinishException);
  }

  @Test
  public void testTaskCancellationTimeout() throws InterruptedException {
    final AtomicReference<Throwable> cancelActionValue = new AtomicReference<>();
    final CountDownLatch runLatch = new CountDownLatch(1);
    final CountDownLatch listenerLatch = new CountDownLatch(1);
    Task<Integer> uncompleted = Task.async(() -> {
      runLatch.countDown();
      return Promises.settable();
    });
    uncompleted.addListener(p -> {
      if (p.isFailed() && Exceptions.isCancellation(p.getError())) {
        cancelActionValue.set(p.getError().getCause());
      }
      listenerLatch.countDown();
    } );

    Task<?> task = uncompleted.withTimeout(10, TimeUnit.MILLISECONDS).recover(e -> 0);
    runAndWait("TestTaskCancellation.testTaskCancellationTimeout", task);
    assertTrue(listenerLatch.await(5, TimeUnit.SECONDS));

    assertTrue(cancelActionValue.get() instanceof EarlyFinishException);
  }

}

@@ -1268,32 +1268,95 @@
 ntDownLatch(1);%0A
+    final CountDownLatch pendingLatch = new CountDownLatch(1);%0A
     Task%3C?%3E unco
@@ -1376,32 +1376,32 @@
 k.async(() -%3E %7B%0A
-
       runLatch.c
@@ -1676,32 +1676,41 @@
 getEngine().run(
+Task.par(
 uncompleted);%0A  
@@ -1708,29 +1708,388 @@
 eted
-);%0A    assertTrue(run
+, Task.action(() -%3E pendingLatch.countDown())));%0A    assertTrue(runLatch.await(5, TimeUnit.SECONDS));%0A    /*%0A     *   pendingLatch makes sure that uncompleted transitioned to PENDING state%0A     *   This is a trick, we are relying here on a fact that Task.par will schedule%0A     *   tasks in order in which they are being passed to Tasp.par() method.%0A     */%0A    assertTrue(pending
 Latc

Remove race from TestTaskCancellation.