core: Alternate ipV4 and ipV6 addresses for Happy Eyeballs in PickFirstLeafLoadBalancer

core/src/main/java/io/grpc/internal/PickFirstLeafLoadBalancer.java

@@ -34,6 +34,8 @@
 import io.grpc.LoadBalancer;
 import io.grpc.Status;
 import io.grpc.SynchronizationContext.ScheduledHandle;
+import java.net.Inet4Address;
+import java.net.InetSocketAddress;
 import java.net.SocketAddress;
 import java.util.ArrayList;
 import java.util.Collections;
@@ -60,7 +62,7 @@
   static final int CONNECTION_DELAY_INTERVAL_MS = 250;
   private final Helper helper;
   private final Map<SocketAddress, SubchannelData> subchannels = new HashMap<>();
-  private final Index addressIndex = new Index(ImmutableList.of());
+  private final Index addressIndex = Index.create(ImmutableList.of());
   private int numTf = 0;
   private boolean firstPass = true;
   @Nullable
@@ -609,26 +611,71 @@
     }
   }
 
+  interface Index {
+    static Index create(List<EquivalentAddressGroup> groups) {
+      if (PickFirstLoadBalancerProvider.isEnabledHappyEyeballs()) {
+        return new IndexHappyEyeballs(groups);
+      } else {
+        return new IndexNonHE(groups);
+      }
+    }
+
+    boolean isValid();
+
+    boolean isAtBeginning();
+
+    /**
+     * Move to next address in group.  If last address in group move to first address of next group.
+     *
+     * @return false if went off end of the list, otherwise true
+     */
+    boolean increment();
+
+    void reset();
+
+    SocketAddress getCurrentAddress();
+
+    Attributes getCurrentEagAttributes();
+
+    List<EquivalentAddressGroup> getCurrentEagAsList();
+
+    /**
+     * Update to new groups, resetting the current index.
+     */
+    void updateGroups(List<EquivalentAddressGroup> newGroups);
+
+    /**
+     * Returns false if the needle was not found and the current index was left unchanged.
+     */
+    boolean seekTo(SocketAddress needle);
+
+    int size();
+
+    int getGroupIndex();
+  }
+
   /**
-   * Index as in 'i', the pointer to an entry. Not a "search index."
+   * IndexNonHE as in 'i', the pointer to an entry. Not a "search index."
    * All updates should be done in a synchronization context.
    */
   @VisibleForTesting
-  static final class Index {
+  private static final class IndexNonHE implements Index {
     private List<EquivalentAddressGroup> addressGroups;
     private int size;
     private int groupIndex;
     private int addressIndex;
 
-    public Index(List<EquivalentAddressGroup> groups) {
+    public IndexNonHE(List<EquivalentAddressGroup> groups) {
       updateGroups(groups);
     }
 
+    @Override
     public boolean isValid() {
       // Is invalid if empty or has incremented off the end
       return groupIndex < addressGroups.size();
     }
 
+    @Override
     public boolean isAtBeginning() {
       return groupIndex == 0 && addressIndex == 0;
     }
@@ -637,6 +684,7 @@
      * Move to next address in group.  If last address in group move to first address of next group.
      * @return false if went off end of the list, otherwise true
      */
+    @Override
     public boolean increment() {
       if (!isValid()) {
         return false;
@@ -653,25 +701,29 @@
       return true;
     }
 
+    @Override
     public void reset() {
       groupIndex = 0;
       addressIndex = 0;
     }
 
+    @Override
     public SocketAddress getCurrentAddress() {
       if (!isValid()) {
         throw new IllegalStateException("Index is past the end of the address group list");
       }
       return addressGroups.get(groupIndex).getAddresses().get(addressIndex);
     }
 
+    @Override
     public Attributes getCurrentEagAttributes() {
       if (!isValid()) {
         throw new IllegalStateException("Index is off the end of the address group list");
       }
       return addressGroups.get(groupIndex).getAttributes();
     }
 
+    @Override
     public List<EquivalentAddressGroup> getCurrentEagAsList() {
       return Collections.singletonList(
           new EquivalentAddressGroup(getCurrentAddress(), getCurrentEagAttributes()));
@@ -680,6 +732,7 @@
     /**
      * Update to new groups, resetting the current index.
      */
+    @Override
     public void updateGroups(List<EquivalentAddressGroup> newGroups) {
       addressGroups = checkNotNull(newGroups, "newGroups");
       reset();
@@ -693,6 +746,7 @@
     /**
      * Returns false if the needle was not found and the current index was left unchanged.
      */
+    @Override
     public boolean seekTo(SocketAddress needle) {
       for (int i = 0; i < addressGroups.size(); i++) {
         EquivalentAddressGroup group = addressGroups.get(i);
@@ -707,14 +761,171 @@
       return false;
     }
 
+    @Override
     public int size() {
       return size;
     }
+
+    @Override
+    public int getGroupIndex() {
+      return groupIndex;
+    }
+  }
+
+  private static final class IndexHappyEyeballs implements Index {
+    private List<EquivalentAddressGroup> addressGroups;
+    private List<InterleavedEntry> interleavedAddresses = new ArrayList<>();
+    private int interleavedIndex = 0;
+
+    public IndexHappyEyeballs(List<EquivalentAddressGroup> groups) {
+      updateGroups(groups);
+    }
+
+    @Override
+    public boolean increment() {
+      if (!isValid()) {
+        return false;
+      }
+
+      interleavedIndex++;
+
+      return isValid();
+    }
+
+    @Override
+    public boolean isValid() {
+      return interleavedIndex < interleavedAddresses.size();
+    }
+
+    @Override
+    public boolean isAtBeginning() {
+      return interleavedIndex == 0;
+    }
+
+    @Override
+    public void reset() {
+      interleavedIndex = 0;
+    }
+
+    @Override
+    public SocketAddress getCurrentAddress() {
+      if (!isValid()) {
+        throw new IllegalStateException("Index is past the end of the address group list");
+      }
+      return interleavedAddresses.get(interleavedIndex).address;
+    }
+
+    @Override
+    public Attributes getCurrentEagAttributes() {
+      return getCurrentEag().getAttributes();
+    }
+
+    @Override
+    public List<EquivalentAddressGroup> getCurrentEagAsList() {
+      return Collections.singletonList(getCurrentEag());
+    }
+
+    @Override
+    public int getGroupIndex() {
+      if (!isValid()) {
+        throw new IllegalStateException("Index is past the end of the address group list");
+      }
+      return interleavedAddresses.get(interleavedIndex).addressGroup;
+    }
+
+    private EquivalentAddressGroup getCurrentEag() {
+      if (!isValid()) {
+        throw new IllegalStateException("Index is past the end of the address group list");
+      }
+      return addressGroups.get(interleavedAddresses.get(interleavedIndex).addressGroup);
+    }
+
+    @Override
+    public void updateGroups(List<EquivalentAddressGroup> newGroups) {
+      addressGroups = checkNotNull(newGroups, "newGroups");
+      Boolean firstIsV6 = null;
+      List<InterleavedEntry> v4Entries = new ArrayList<>();
+      List<InterleavedEntry> v6Entries = new ArrayList<>();
+      for (int g = 0; g <  newGroups.size(); g++) {
+        EquivalentAddressGroup eag = newGroups.get(g);
+        for (int a = 0; a < eag.getAddresses().size(); a++) {
+          SocketAddress addr = eag.getAddresses().get(a);
+          boolean isIpV4 = addr instanceof InetSocketAddress
+              && ((InetSocketAddress) addr).getAddress() instanceof Inet4Address;
+          if (isIpV4) {
+            if (firstIsV6 == null) {
+              firstIsV6 = false;
+            }
+            v4Entries.add(new InterleavedEntry(g, addr));
+          } else {
+            if (firstIsV6 == null) {
+              firstIsV6 = true;
+            }
+            v6Entries.add(new InterleavedEntry(g, addr));
+          }
+        }
+      }
+
+      this.interleavedAddresses =
+          firstIsV6 != null && firstIsV6
+          ? interleave(v6Entries, v4Entries)
+          : interleave(v4Entries, v6Entries);
+
+      reset();
+    }
+
+    private static List<InterleavedEntry> interleave(List<InterleavedEntry> firstFamily,
+                                                     List<InterleavedEntry> secondFamily) {
+      if (firstFamily.isEmpty()) {
+        return secondFamily;
+      }
+      if (secondFamily.isEmpty()) {
+        return firstFamily;
+      }
+
+      List<InterleavedEntry> result = new ArrayList<>();
+      for (int i = 0; i < Math.max(firstFamily.size(), secondFamily.size()); i++) {
+        if (i < firstFamily.size()) {
+          result.add(firstFamily.get(i));
+        }
+        if (i < secondFamily.size()) {
+          result.add(secondFamily.get(i));
+        }
+      }
+      return result;
+    }
+
+    @Override
+    public boolean seekTo(SocketAddress needle) {
+      checkNotNull(needle, "needle");
+      for (int i = 0; i < interleavedAddresses.size(); i++) {
+        if (interleavedAddresses.get(i).address.equals(needle)) {
+          this.interleavedIndex = i;
+          return true;
+        }
+      }
+      return false;
+    }
+
+    @Override
+    public int size() {
+      return interleavedAddresses.size();
+    }
+
+    private static final class InterleavedEntry {
+      private final int addressGroup;
+      private final SocketAddress address;
+
+      public InterleavedEntry(int addressGroup, SocketAddress address) {
+        this.addressGroup = addressGroup;
+        this.address = address;
+      }
+    }
   }
 
   @VisibleForTesting
   int getGroupIndex() {
-    return addressIndex.groupIndex;
+    return addressIndex.getGroupIndex();
   }
 
   @VisibleForTesting
@@ -778,4 +989,5 @@
       this.randomSeed = randomSeed;
     }
   }
+
 }