Merge pull request #1296 from firstway/allow_same_ports_with_differen…

…t_protocols allow same static external ports with diff protocols
spotify · May 7, 2020 · 60d1212 · 60d1212
2 parents 535adc7 + d1dc2de
commit 60d1212
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 12 deletions.
diff --git a/helios-services/src/main/java/com/spotify/helios/agent/PortAllocator.java b/helios-services/src/main/java/com/spotify/helios/agent/PortAllocator.java
@@ -26,6 +26,7 @@
 import java.io.IOException;
 import java.net.ServerSocket;
 import java.util.Collections;
+import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
@@ -69,10 +70,11 @@ public Map<String, Integer> allocate(final Map<String, PortMapping> ports,
     return allocate0(ports, Sets.newHashSet(used));
   }
 
-  private Map<String, Integer> allocate0(final Map<String, PortMapping> mappings,
-                                         final Set<Integer> used) {
+  private Map<String, Integer> allocate0(
+      final Map<String, PortMapping> mappings, final Set<Integer> used) {
 
     final ImmutableMap.Builder<String, Integer> allocation = ImmutableMap.builder();
+    final Set<Integer> staticExternalPorts = new HashSet<>();
 
     for (final Map.Entry<String, PortMapping> entry : mappings.entrySet()) {
       final String name = entry.getKey();
@@ -84,27 +86,43 @@ private Map<String, Integer> allocate0(final Map<String, PortMapping> mappings,
           return null;
         }
       } else {
-        if (!allocateStatic(allocation, used, name, externalPort)) {
-          return null;
-        }
+        // for static external ports, it's valid to
+        // have two same ports with different protocols
+        // e.g.
+        // "externalPort": 11101, "protocol": "tcp"
+        // "externalPort": 11101, "protocol": "udp"
+        staticExternalPorts.add(externalPort);
+      }
+    }
+
+    // check conflict between `used` and `staticExternalPorts`
+    for (final Integer externalPort : staticExternalPorts) {
+      if (used.contains(externalPort)) {
+        return null;
+      }
+    }
+
+    // we have checked the conflict between `used` and `staticExternalPorts` above
+    // so here we can add all externalPorts into `allocation` without any validation
+    for (final Map.Entry<String, PortMapping> entry : mappings.entrySet()) {
+      final String name = entry.getKey();
+      final PortMapping portMapping = entry.getValue();
+      final Integer externalPort = portMapping.getExternalPort();
+
+      if (externalPort != null) {
+        allocateStatic(allocation, used, name, externalPort);
       }
     }
 
     return allocation.build();
   }
 
-  private boolean allocateStatic(final ImmutableMap.Builder<String, Integer> allocation,
+  private void allocateStatic(final ImmutableMap.Builder<String, Integer> allocation,
                                  final Set<Integer> used,
                                  final String name,
                                  final Integer port) {
-    // Verify that this port is not in use
-    if (used.contains(port)) {
-      return false;
-    }
-
     used.add(port);
     allocation.put(name, port);
-    return true;
   }
 
   private boolean allocateDynamic(final ImmutableMap.Builder<String, Integer> allocation,

diff --git a/helios-services/src/test/java/com/spotify/helios/agent/PortAllocatorTest.java b/helios-services/src/test/java/com/spotify/helios/agent/PortAllocatorTest.java
@@ -49,6 +49,23 @@ public void testAllocate() throws Exception {
     assertThat(allocation, hasEntry("p2", 18));
   }
 
+  @Test
+  public void testAllocatePortsWithDifferentProtocols() throws Exception {
+    final PortAllocator sut = new PortAllocator(20000, 20010);
+    final Map<String, PortMapping> mapping =
+        ImmutableMap.of(
+            "p1", PortMapping.of(17),
+            "p-tcp", PortMapping.of(18, 18),
+            "p-udp", PortMapping.of(18, 18));
+    final Set<Integer> used = ImmutableSet.of(10, 11);
+    final Map<String, Integer> allocation = sut.allocate(mapping, used);
+    assertThat(
+        allocation,
+        hasEntry(is("p1"), allOf(greaterThanOrEqualTo(20000), lessThanOrEqualTo(20010))));
+    assertThat(allocation, hasEntry("p-tcp", 18));
+    assertThat(allocation, hasEntry("p-udp", 18));
+  }
+
   @Test
   public void testInsufficientPortsFail1() throws Exception {
     final PortAllocator sut = new PortAllocator(10, 11);