Fix a race condition in CompositeEndpointGroup

core/src/main/java/com/linecorp/armeria/client/endpoint/CompositeEndpointGroup.java

@@ -20,7 +20,7 @@
 import java.util.List;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ScheduledExecutorService;
-import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicReference;
 
 import com.google.common.base.MoreObjects;
 import com.google.common.collect.ImmutableList;
@@ -40,31 +40,26 @@ final class CompositeEndpointGroup extends AbstractEndpointGroup implements List
     private final List<EndpointGroup> endpointGroups;
 
     private final CompletableFuture<List<Endpoint>> initialEndpointsFuture;
-    private final AtomicBoolean dirty;
 
     private final EndpointSelectionStrategy selectionStrategy;
     private final EndpointSelector selector;
     private final long selectionTimeoutMillis;
 
     private final AsyncCloseableSupport closeable = AsyncCloseableSupport.of(this::closeAsync);
 
-    private volatile List<Endpoint> merged = ImmutableList.of();
+    private final AtomicReference<List<Endpoint>> merged = new AtomicReference<>();
 
     /**
      * Constructs a new {@link CompositeEndpointGroup} that merges all the given {@code endpointGroups}.
      */
     CompositeEndpointGroup(EndpointSelectionStrategy selectionStrategy,
                            Iterable<EndpointGroup> endpointGroups) {
-
         this.endpointGroups = ImmutableList.copyOf(requireNonNull(endpointGroups, "endpointGroups"));
-        dirty = new AtomicBoolean(true);
+        merged.set(ImmutableList.of());
 
         long selectionTimeoutMillis = 0;
         for (EndpointGroup endpointGroup : endpointGroups) {
-            endpointGroup.addListener(unused -> {
-                dirty.set(true);
-                notifyListeners(endpoints());
-            });
+            endpointGroup.addListener(unused -> notifyListeners(rebuildEndpoints()));
             selectionTimeoutMillis = Math.max(selectionTimeoutMillis,
                                               endpointGroup.selectionTimeoutMillis());
         }
@@ -74,30 +69,33 @@ final class CompositeEndpointGroup extends AbstractEndpointGroup implements List
                 CompletableFuture.anyOf(this.endpointGroups.stream()
                                                            .map(EndpointGroup::whenReady)
                                                            .toArray(CompletableFuture[]::new))
-                                 .thenApply(unused -> endpoints());
+                                 .thenApply(unused -> rebuildEndpoints());
 
         this.selectionStrategy = requireNonNull(selectionStrategy, "selectionStrategy");
         selector = requireNonNull(selectionStrategy, "selectionStrategy").newSelector(this);
     }
 
     @Override
     public List<Endpoint> endpoints() {
-        if (!dirty.get()) {
-            return merged;
-        }
-
-        if (!dirty.compareAndSet(true, false)) {
-            // Another thread might be updating merged at this time, but endpoint groups are allowed to take a
-            // little bit of time to reflect updates.
-            return merged;
-        }
+        final List<Endpoint> endpoints = merged.get();
+        assert endpoints != null;
+        return endpoints;
+    }
 
-        final ImmutableList.Builder<Endpoint> newEndpoints = ImmutableList.builder();
-        for (EndpointGroup endpointGroup : endpointGroups) {
-            newEndpoints.addAll(endpointGroup.endpoints());
+    private List<Endpoint> rebuildEndpoints() {
+        for (;;) {
+            // Get the current endpoints before making a new endpoints list.
+            final List<Endpoint> oldEndpoints = merged.get();
+            final ImmutableList.Builder<Endpoint> newEndpointsBuilder = ImmutableList.builder();
+            for (EndpointGroup endpointGroup : endpointGroups) {
+                newEndpointsBuilder.addAll(endpointGroup.endpoints());
+            }
+            final List<Endpoint> newEndpoints = newEndpointsBuilder.build();
+            if (merged.compareAndSet(oldEndpoints, newEndpoints)) {
+                return newEndpoints;
+            }
+            // Changed by another thread while we were building a new one. Try again.
         }
-
-        return merged = newEndpoints.build();
     }
 
     @Override