- // This method could be synchronized, but that would mean that concurrent attempts to load an invalid augmentation
- // would end up being unnecessarily contended -- blocking real progress and not being able to run concurrently
- // while producing no effect. We therefore use optimistic read + CAS.
- private void reloadAllAugmentations() {
- // Load current values
- Augmentations oldAugmentations = augmentations;
-
- // FIXME: can we detect when we have both maps fully populated and skip all of this?
-
- // Scratch space for additions
- final Map<PathArgument, DataContainerCodecPrototype<?>> addByYang = new HashMap<>();
- final Map<Class<?>, DataContainerCodecPrototype<?>> addByStream = new HashMap<>();
-
- // Iterate over all possibilities, checking for modifications.
- for (final Type augment : possibleAugmentations.values()) {
- final DataContainerCodecPrototype<?> augProto = getAugmentationPrototype(augment);
- if (augProto != null) {
- final PathArgument yangArg = augProto.getYangArg();
- final Class<?> bindingClass = augProto.getBindingClass();
- if (!oldAugmentations.byYang.containsKey(yangArg)) {
- if (addByYang.putIfAbsent(yangArg, augProto) == null) {
- LOG.trace("Discovered new YANG mapping {} -> {} in {}", yangArg, augProto, this);
- }
- }
- if (!oldAugmentations.byStream.containsKey(bindingClass)) {
- if (addByStream.putIfAbsent(bindingClass, augProto) == null) {
- LOG.trace("Discovered new class mapping {} -> {} in {}", bindingClass, augProto, this);
- }
- }
- }
- }
-
- while (true) {
- if (addByYang.isEmpty() && addByStream.isEmpty()) {
- LOG.trace("No new augmentations discovered in {}", this);
- return;
- }
-
- // We have some additions, propagate them out
- final Augmentations newAugmentations = new Augmentations(concatMaps(oldAugmentations.byYang, addByYang),
- concatMaps(oldAugmentations.byStream, addByStream));
- if (AUGMENTATIONS_UPDATER.compareAndSet(this, oldAugmentations, newAugmentations)) {
- // Success, we are done
- return;
- }
-
- // We have raced installing new augmentations, read them again, remove everything present in the installed
- // once and try again. This may mean that we end up not doing anything, but that's fine.
- oldAugmentations = augmentations;
-
- // We could use Map.removeAll(oldAugmentations.byYang.keySet()), but that forces the augmentation's keyset
- // to be materialized, which we otherwise do not need. Hence we do this the other way around, instantiating
- // our temporary maps' keySets and iterating over them. That's fine as we'll be throwing those maps away.
- removeMapKeys(addByYang, oldAugmentations.byYang);
- removeMapKeys(addByStream, oldAugmentations.byStream);
- }
- }
-
- private static <K, V> ImmutableMap<K, V> concatMaps(final ImmutableMap<K, V> old, final Map<K, V> add) {
- if (add.isEmpty()) {
- return old;
- }
-
- final Builder<K, V> builder = ImmutableMap.builderWithExpectedSize(old.size() + add.size());
- builder.putAll(old);
- builder.putAll(add);
- return builder.build();
- }
-
- private static <K, V> void removeMapKeys(final Map<K, V> removeFrom, final ImmutableMap<K, V> map) {
- final Iterator<K> it = removeFrom.keySet().iterator();
- while (it.hasNext()) {
- if (map.containsKey(it.next())) {
- it.remove();
- }
- }