[SPARK-27726][CORE] Fix performance of ElementTrackingStore deletes when using InMemoryStore under high loads

The details of the PR are explored in-depth in the sub-tasks of the umbrella jira SPARK-27726.
Briefly:
  1. Stop issuing asynchronous requests to cleanup elements in the tracking store when a request is already pending
  2. Fix a couple of thread-safety issues (mutable state and mis-ordered updates)
  3. Move Summary deletion outside of Stage deletion loop like Tasks already are
  4. Reimplement multi-delete in a removeAllKeys call which allows InMemoryStore to implement it in a performant manner.
  5. Some generic typing and exception handling cleanup

We see about five orders of magnitude improvement in the deletion code, which for us is the difference between a server that needs restarting daily, and one that is stable over weeks.

Unit tests for the fire-once asynchronous code and the removeAll calls in both LevelDB and InMemoryStore are supplied.  It was noted that the testing code for the LevelDB and InMemoryStore is highly repetitive, and should probably be merged, but we did not attempt that in this PR.

A version of this code was run in our production 2.3.3 and we were able to sustain higher throughput without going into GC overload (which was happening on a daily basis some weeks ago).

A version of this code was also put under a purpose-built Performance Suite of tests to verify performance under both types of Store implementations for both before and after code streams and for both total and partial delete cases (this code is not included in this PR).

Closes apache#24616 from davidnavas/PentaBugFix.

Authored-by: David Navas <[email protected]>
Signed-off-by: Marcelo Vanzin <[email protected]>

Author: David Navas

common/kvstore/src/main/java/org/apache/spark/util/kvstore/InMemoryStore.java

@@ -21,16 +21,18 @@
 import java.util.Collection;
 import java.util.Collections;
 import java.util.Iterator;
+import java.util.HashSet;
 import java.util.List;
 import java.util.NoSuchElementException;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
+import java.util.function.BiConsumer;
+import java.util.function.Predicate;
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
 import com.google.common.base.Objects;
 import com.google.common.base.Preconditions;
-import com.google.common.base.Throwables;
 
 import org.apache.spark.annotation.Private;
 
@@ -43,7 +45,7 @@
 public class InMemoryStore implements KVStore {
 
   private Object metadata;
-  private ConcurrentMap<Class<?>, InstanceList> data = new ConcurrentHashMap<>();
+  private InMemoryLists inMemoryLists = new InMemoryLists();
 
   @Override
   public <T> T getMetadata(Class<T> klass) {
@@ -57,13 +59,13 @@ public void setMetadata(Object value) {
 
   @Override
   public long count(Class<?> type) {
-    InstanceList list = data.get(type);
+    InstanceList<?> list = inMemoryLists.get(type);
     return list != null ? list.size() : 0;
   }
 
   @Override
   public long count(Class<?> type, String index, Object indexedValue) throws Exception {
-    InstanceList list = data.get(type);
+    InstanceList<?> list = inMemoryLists.get(type);
     int count = 0;
     Object comparable = asKey(indexedValue);
     KVTypeInfo.Accessor accessor = list.getIndexAccessor(index);
@@ -77,45 +79,51 @@ public long count(Class<?> type, String index, Object indexedValue) throws Excep
 
   @Override
   public <T> T read(Class<T> klass, Object naturalKey) {
-    InstanceList list = data.get(klass);
-    Object value = list != null ? list.get(naturalKey) : null;
+    InstanceList<T> list = inMemoryLists.get(klass);
+    T value = list != null ? list.get(naturalKey) : null;
     if (value == null) {
       throw new NoSuchElementException();
     }
-    return klass.cast(value);
+    return value;
   }
 
   @Override
   public void write(Object value) throws Exception {
-    InstanceList list = data.computeIfAbsent(value.getClass(), key -> {
-      try {
-        return new InstanceList(key);
-      } catch (Exception e) {
-        throw Throwables.propagate(e);
-      }
-    });
-    list.put(value);
+    inMemoryLists.write(value);
   }
 
   @Override
   public void delete(Class<?> type, Object naturalKey) {
-    InstanceList list = data.get(type);
+    InstanceList<?> list = inMemoryLists.get(type);
     if (list != null) {
       list.delete(naturalKey);
     }
   }
 
   @Override
   public <T> KVStoreView<T> view(Class<T> type){
-    InstanceList list = data.get(type);
-    return list != null ? list.view(type)
-      : new InMemoryView<>(type, Collections.<T>emptyList(), null);
+    InstanceList<T> list = inMemoryLists.get(type);
+    return list != null ? list.view() : emptyView();
   }
 
   @Override
   public void close() {
     metadata = null;
-    data.clear();
+    inMemoryLists.clear();
+  }
+
+  @Override
+  public <T> boolean removeAllByIndexValues(
+      Class<T> klass,
+      String index,
+      Collection<?> indexValues) {
+    InstanceList<T> list = inMemoryLists.get(klass);
+
+    if (list != null) {
+      return list.countingRemoveAllByIndexValues(index, indexValues) > 0;
+    } else {
+      return false;
+    }
   }
 
   @SuppressWarnings("unchecked")
@@ -126,64 +134,150 @@ private static Comparable<Object> asKey(Object in) {
     return (Comparable<Object>) in;
   }
 
-  private static class InstanceList {
+  @SuppressWarnings("unchecked")
+  private static <T> KVStoreView<T> emptyView() {
+    return (InMemoryView<T>) InMemoryView.EMPTY_VIEW;
+  }
+
+  /**
+   * Encapsulates ConcurrentHashMap so that the typing in and out of the map strictly maps a
+   * class of type T to an InstanceList of type T.
+   */
+  private static class InMemoryLists {
+    private final ConcurrentMap<Class<?>, InstanceList<?>> data = new ConcurrentHashMap<>();
+
+    @SuppressWarnings("unchecked")
+    public <T> InstanceList<T> get(Class<T> type) {
+      return (InstanceList<T>) data.get(type);
+    }
+
+    @SuppressWarnings("unchecked")
+    public <T> void write(T value) throws Exception {
+      InstanceList<T> list =
+        (InstanceList<T>) data.computeIfAbsent(value.getClass(), InstanceList::new);
+      list.put(value);
+    }
+
+    public void clear() {
+      data.clear();
+    }
+  }
+
+  private static class InstanceList<T> {
+
+    /**
+     * A BiConsumer to control multi-entity removal.  We use this in a forEach rather than an
+     * iterator because there is a bug in jdk8 which affects remove() on all concurrent map
+     * iterators.  https://bugs.openjdk.java.net/browse/JDK-8078645
+     */
+    private static class CountingRemoveIfForEach<T> implements BiConsumer<Comparable<Object>, T> {
+      private final ConcurrentMap<Comparable<Object>, T> data;
+      private final Predicate<? super T> filter;
+
+      /**
+       * Keeps a count of the number of elements removed.  This count is not currently surfaced
+       * to clients of KVStore as Java's generic removeAll() construct returns only a boolean,
+       * but I found it handy to have the count of elements removed while debugging; a count being
+       * no more complicated than a boolean, I've retained that behavior here, even though there
+       * is no current requirement.
+       */
+      private int count = 0;
+
+      CountingRemoveIfForEach(
+          ConcurrentMap<Comparable<Object>, T> data,
+          Predicate<? super T> filter) {
+        this.data = data;
+        this.filter = filter;
+      }
+
+      @Override
+      public void accept(Comparable<Object> key, T value) {
+        if (filter.test(value)) {
+          if (data.remove(key, value)) {
+            count++;
+          }
+        }
+      }
+
+      public int count() { return count; }
+    }
 
     private final KVTypeInfo ti;
     private final KVTypeInfo.Accessor naturalKey;
-    private final ConcurrentMap<Comparable<Object>, Object> data;
-
-    private int size;
+    private final ConcurrentMap<Comparable<Object>, T> data;
 
-    private InstanceList(Class<?> type) throws Exception {
-      this.ti = new KVTypeInfo(type);
+    private InstanceList(Class<?> klass) {
+      this.ti = new KVTypeInfo(klass);
       this.naturalKey = ti.getAccessor(KVIndex.NATURAL_INDEX_NAME);
       this.data = new ConcurrentHashMap<>();
-      this.size = 0;
     }
 
     KVTypeInfo.Accessor getIndexAccessor(String indexName) {
       return ti.getAccessor(indexName);
     }
 
-    public Object get(Object key) {
+    int countingRemoveAllByIndexValues(String index, Collection<?> indexValues) {
+      Predicate<? super T> filter = getPredicate(ti.getAccessor(index), indexValues);
+      CountingRemoveIfForEach<T> callback = new CountingRemoveIfForEach<>(data, filter);
+
+      data.forEach(callback);
+      return callback.count();
+    }
+
+    public T get(Object key) {
       return data.get(asKey(key));
     }
 
-    public void put(Object value) throws Exception {
-      Preconditions.checkArgument(ti.type().equals(value.getClass()),
-        "Unexpected type: %s", value.getClass());
-      if (data.put(asKey(naturalKey.get(value)), value) == null) {
-        size++;
-      }
+    public void put(T value) throws Exception {
+      data.put(asKey(naturalKey.get(value)), value);
     }
 
     public void delete(Object key) {
-      if (data.remove(asKey(key)) != null) {
-        size--;
-      }
+      data.remove(asKey(key));
     }
 
     public int size() {
-      return size;
+      return data.size();
     }
 
-    @SuppressWarnings("unchecked")
-    public <T> InMemoryView<T> view(Class<T> type) {
-      Preconditions.checkArgument(ti.type().equals(type), "Unexpected type: %s", type);
-      Collection<T> all = (Collection<T>) data.values();
-      return new InMemoryView<>(type, all, ti);
+    public InMemoryView<T> view() {
+      return new InMemoryView<>(data.values(), ti);
+    }
+
+    private static <T> Predicate<? super T> getPredicate(
+        KVTypeInfo.Accessor getter,
+        Collection<?> values) {
+      if (Comparable.class.isAssignableFrom(getter.getType())) {
+        HashSet<?> set = new HashSet<>(values);
+
+        return (value) -> set.contains(indexValueForEntity(getter, value));
+      } else {
+        HashSet<Comparable> set = new HashSet<>(values.size());
+        for (Object key : values) {
+          set.add(asKey(key));
+        }
+        return (value) -> set.contains(asKey(indexValueForEntity(getter, value)));
+      }
     }
 
+    private static Object indexValueForEntity(KVTypeInfo.Accessor getter, Object entity) {
+      try {
+        return getter.get(entity);
+      } catch (ReflectiveOperationException e) {
+        throw new RuntimeException(e);
+      }
+    }
   }
 
   private static class InMemoryView<T> extends KVStoreView<T> {
+    private static final InMemoryView<?> EMPTY_VIEW =
+      new InMemoryView<>(Collections.emptyList(), null);
 
     private final Collection<T> elements;
     private final KVTypeInfo ti;
     private final KVTypeInfo.Accessor natural;
 
-    InMemoryView(Class<T> type, Collection<T> elements, KVTypeInfo ti) {
-      super(type);
+    InMemoryView(Collection<T> elements, KVTypeInfo ti) {
       this.elements = elements;
       this.ti = ti;
       this.natural = ti != null ? ti.getAccessor(KVIndex.NATURAL_INDEX_NAME) : null;
@@ -195,34 +289,32 @@ public Iterator<T> iterator() {
         return new InMemoryIterator<>(elements.iterator());
       }
 
-      try {
-        KVTypeInfo.Accessor getter = index != null ? ti.getAccessor(index) : null;
-        int modifier = ascending ? 1 : -1;
+      KVTypeInfo.Accessor getter = index != null ? ti.getAccessor(index) : null;
+      int modifier = ascending ? 1 : -1;
 
-        final List<T> sorted = copyElements();
-        Collections.sort(sorted, (e1, e2) -> modifier * compare(e1, e2, getter));
-        Stream<T> stream = sorted.stream();
+      final List<T> sorted = copyElements();
+      sorted.sort((e1, e2) -> modifier * compare(e1, e2, getter));
+      Stream<T> stream = sorted.stream();
 
-        if (first != null) {
-          stream = stream.filter(e -> modifier * compare(e, getter, first) >= 0);
-        }
-
-        if (last != null) {
-          stream = stream.filter(e -> modifier * compare(e, getter, last) <= 0);
-        }
+      if (first != null) {
+        Comparable<?> firstKey = asKey(first);
+        stream = stream.filter(e -> modifier * compare(e, getter, firstKey) >= 0);
+      }
 
-        if (skip > 0) {
-          stream = stream.skip(skip);
-        }
+      if (last != null) {
+        Comparable<?> lastKey = asKey(last);
+        stream = stream.filter(e -> modifier * compare(e, getter, lastKey) <= 0);
+      }
 
-        if (max < sorted.size()) {
-          stream = stream.limit((int) max);
-        }
+      if (skip > 0) {
+        stream = stream.skip(skip);
+      }
 
-        return new InMemoryIterator<>(stream.iterator());
-      } catch (Exception e) {
-        throw Throwables.propagate(e);
+      if (max < sorted.size()) {
+        stream = stream.limit((int) max);
       }
+
+      return new InMemoryIterator<>(stream.iterator());
     }
 
     /**
@@ -232,9 +324,10 @@ private List<T> copyElements() {
       if (parent != null) {
         KVTypeInfo.Accessor parentGetter = ti.getParentAccessor(index);
         Preconditions.checkArgument(parentGetter != null, "Parent filter for non-child index.");
+        Comparable<?> parentKey = asKey(parent);
 
         return elements.stream()
-          .filter(e -> compare(e, parentGetter, parent) == 0)
+          .filter(e -> compare(e, parentGetter, parentKey) == 0)
           .collect(Collectors.toList());
       } else {
         return new ArrayList<>(elements);
@@ -243,24 +336,23 @@ private List<T> copyElements() {
 
     private int compare(T e1, T e2, KVTypeInfo.Accessor getter) {
       try {
-        int diff = compare(e1, getter, getter.get(e2));
+        int diff = compare(e1, getter, asKey(getter.get(e2)));
         if (diff == 0 && getter != natural) {
-          diff = compare(e1, natural, natural.get(e2));
+          diff = compare(e1, natural, asKey(natural.get(e2)));
         }
         return diff;
-      } catch (Exception e) {
-        throw Throwables.propagate(e);
+      } catch (ReflectiveOperationException e) {
+        throw new RuntimeException(e);
       }
     }
 
-    private int compare(T e1, KVTypeInfo.Accessor getter, Object v2) {
+    private int compare(T e1, KVTypeInfo.Accessor getter, Comparable<?> v2) {
       try {
-        return asKey(getter.get(e1)).compareTo(asKey(v2));
-      } catch (Exception e) {
-        throw Throwables.propagate(e);
+        return asKey(getter.get(e1)).compareTo(v2);
+      } catch (ReflectiveOperationException e) {
+        throw new RuntimeException(e);
       }
     }
-
   }
 
   private static class InMemoryIterator<T> implements KVStoreIterator<T> {