Merge pull request #1029 from SpiNNakerManchester/safe-weak

Make EpochMap safe against concurrent updates by the GC

SpiNNaker-allocserv/src/main/java/uk/ac/manchester/spinnaker/alloc/allocator/Epochs.java

@@ -15,18 +15,19 @@
  */
 package uk.ac.manchester.spinnaker.alloc.allocator;
 
-import static java.util.Collections.newSetFromMap;
 import static java.util.Objects.nonNull;
 import static org.slf4j.LoggerFactory.getLogger;
 
 import java.time.Duration;
 import java.util.Collection;
+import java.util.ConcurrentModificationException;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.WeakHashMap;
+import java.util.function.BiConsumer;
 
 import javax.annotation.PostConstruct;
 
@@ -35,6 +36,8 @@
 import org.springframework.scheduling.TaskScheduler;
 import org.springframework.stereotype.Service;
 
+import uk.ac.manchester.spinnaker.utils.UsedInJavadocOnly;
+
 /**
  * Manages waiting for values.
  *
@@ -258,11 +261,11 @@ public boolean isValid() {
  * @author Donal Fellows
  */
 class EpochMap {
-	/*
-	 * This is a wrapper around the map to provide exactly the ops we want
-	 * (including correct synchronization), and just them.
-	 */
-	private final Map<Integer, Set<Epochs.Epoch>> map = new HashMap<>();
+	/** The value in {@link #map} leaves. Shared. Unimportant. */
+	private static final Object OBJ = new Object();
+
+	/** A map from integers to weak sets of epochs. */
+	private final Map<Integer, Map<Epochs.Epoch, Object>> map = new HashMap<>();
 
 	synchronized void checkEmptyValues() {
 		map.entrySet().removeIf(entry -> entry.getValue().isEmpty());
@@ -277,14 +280,36 @@ void changed(int id) {
 		}
 	}
 
-	private synchronized Set<Epochs.Epoch> removeSet(int id) {
-		return map.remove(id);
+	/**
+	 * Take the set of epochs for a particular ID.
+	 * <p>
+	 * <strong>CRITICAL TRICKY IMPLEMENTATION POINT!</strong>
+	 * {@link WeakHashMap#forEach(BiConsumer)} does not throw
+	 * {@link ConcurrentModificationException} as it is <em>not</em>
+	 * iterator-based. All the other ways of going through the contents can
+	 * throw if the GC removes an element behind the scenes.
+	 *
+	 * @param id
+	 *            The key into the map.
+	 * @return The removed set of epochs. Empty if the key is absent.
+	 */
+	@UsedInJavadocOnly({
+		BiConsumer.class, ConcurrentModificationException.class
+	})
+	private synchronized Set<Epochs.Epoch> removeSet(Integer id) {
+		var weakmap = map.remove(id);
+		if (weakmap == null) {
+			return null;
+		}
+		// Set.copyOf would be great, but can blow up :-(
+		var result = new HashSet<Epochs.Epoch>();
+		weakmap.forEach((epoch, __) -> result.add(epoch));
+		return result;
 	}
 
 	synchronized void addAll(Epochs.Epoch epoch, List<Integer> ids) {
 		for (var id : ids) {
-			map.computeIfAbsent(id, key -> newSetFromMap(new WeakHashMap<>()))
-					.add(epoch);
+			map.computeIfAbsent(id, key -> new WeakHashMap<>()).put(epoch, OBJ);
 		}
 	}