/*-
 * -\-\-
 * Flo Workflow Definition
 * --
 * Copyright (C) 2016 - 2017 Spotify AB
 * --
 * 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.spotify.flo.context;

import com.spotify.flo.EvalContext;
import com.spotify.flo.Task;
import com.spotify.flo.TaskId;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * A flo {@link EvalContext} that memoizes the results of task evaluations and
 * ensures that tasks are only evaluated once.
 */
public class MemoizingContext extends ForwardingEvalContext {

  private final ConcurrentMap<TaskId, CompletableFuture<Value<?>>> ongoing = new ConcurrentHashMap<>();

  private MemoizingContext(EvalContext baseContext) {
    super(baseContext);
  }

  public static EvalContext composeWith(EvalContext baseContext) {
    return new MemoizingContext(baseContext);
  }

  @SuppressWarnings("unchecked")
  @Override
  public <T> Value<T> evaluateInternal(Task<T> task, EvalContext context) {
    final CompletableFuture<Value<?>> f = new CompletableFuture<>();
    final CompletableFuture<Value<?>> existing = ongoing.putIfAbsent(task.id(), f);
    if (existing != null) {
      return (Value<T>) existing.join();
    }
    final Value<T> value = super.evaluateInternal(task, context);
    f.complete(value);
    return value;
  }
}

@@ -791,55 +791,8 @@
 Id;%0A
-import java.util.concurrent.CompletableFuture;%0A
 impo
@@ -1117,36 +1117,19 @@
 Id, 
-CompletableFuture%3CValu
+Promis
 e%3C?%3E%3E
-%3E
  ong
@@ -1493,110 +1493,456 @@
     
-final CompletableFuture%3CValue%3C?%3E%3E f = new CompletableFuture%3C%3E();%0A    final CompletableFuture%3CValu
+// Unfortunately we cannot use computeIfAbsent here because modification of the CHM%0A    // in computeIfAbsents is not allowed: %22... the computation should be short and simple,%0A    // and must not attempt to update any other mappings of this map.%22.%0A    // However, we could potentially use computeIfAbsent if we rewrite flo task evaluation%0A    // to be iterative instead of recursive.%0A    final Promise%3CT%3E f = context.promise();%0A    final Promis
 e%3C?%3E
-%3E
  exi
@@ -2049,12 +2049,13 @@
 ing.
-join
+value
 ();%0A
@@ -2060,16 +2060,28 @@
 ;%0A    %7D%0A
+    try %7B%0A  
     fina
@@ -2146,27 +2146,113 @@
     
+ 
-f.complete(value
+ value.onFail(f::fail);%0A      value.consume(f::set);%0A    %7D catch (Throwable t) %7B%0A      f.fail(t
 );%0A
+    %7D%0A
     
@@ -2258,21 +2258,25 @@
  return 
+f.
 value
+()
 ;%0A  %7D%0A%7D%0A

use promise & handle exceptions