diff --git a/SpiNNaker-allocserv/src/main/java/uk/ac/manchester/spinnaker/alloc/allocator/Epochs.java b/SpiNNaker-allocserv/src/main/java/uk/ac/manchester/spinnaker/alloc/allocator/Epochs.java index 6810eb4fb3..a92cf6e360 100644 --- a/SpiNNaker-allocserv/src/main/java/uk/ac/manchester/spinnaker/alloc/allocator/Epochs.java +++ b/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> 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> map = new HashMap<>(); synchronized void checkEmptyValues() { map.entrySet().removeIf(entry -> entry.getValue().isEmpty()); @@ -277,14 +280,36 @@ void changed(int id) { } } - private synchronized Set removeSet(int id) { - return map.remove(id); + /** + * Take the set of epochs for a particular ID. + *

+ * CRITICAL TRICKY IMPLEMENTATION POINT! + * {@link WeakHashMap#forEach(BiConsumer)} does not throw + * {@link ConcurrentModificationException} as it is not + * 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 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(); + weakmap.forEach((epoch, __) -> result.add(epoch)); + return result; } synchronized void addAll(Epochs.Epoch epoch, List ids) { for (var id : ids) { - map.computeIfAbsent(id, key -> newSetFromMap(new WeakHashMap<>())) - .add(epoch); + map.computeIfAbsent(id, key -> new WeakHashMap<>()).put(epoch, OBJ); } }