diff --git a/src/main/java/io/github/hectorvent/floci/core/common/port/PortAllocator.java b/src/main/java/io/github/hectorvent/floci/core/common/port/PortAllocator.java index ec31cfd999..c5fc7f00af 100644 --- a/src/main/java/io/github/hectorvent/floci/core/common/port/PortAllocator.java +++ b/src/main/java/io/github/hectorvent/floci/core/common/port/PortAllocator.java @@ -4,34 +4,43 @@ import jakarta.enterprise.context.ApplicationScoped; import jakarta.inject.Inject; -import java.util.concurrent.atomic.AtomicInteger; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; /** - * Thread-safe sequential port dispenser for Lambda Runtime API servers. + * Hands out ports from a configured range for Lambda Runtime API servers. + * Throws {@link IllegalStateException} when the range is exhausted; the + * caller releases a port back to the pool via {@link #release(int)}. */ @ApplicationScoped public class PortAllocator { - private final AtomicInteger counter; private final int basePort; - private final int range; + private final int maxPort; + private final Set inUse = ConcurrentHashMap.newKeySet(); @Inject public PortAllocator(EmulatorConfig config) { - this.basePort = config.services().lambda().runtimeApiBasePort(); - int maxPort = config.services().lambda().runtimeApiMaxPort(); - this.range = maxPort - basePort + 1; - this.counter = new AtomicInteger(0); + this(config.services().lambda().runtimeApiBasePort(), + config.services().lambda().runtimeApiMaxPort()); } PortAllocator(int basePort, int maxPort) { this.basePort = basePort; - this.range = maxPort - basePort + 1; - this.counter = new AtomicInteger(0); + this.maxPort = maxPort; } public int allocate() { - int offset = counter.getAndIncrement(); - return basePort + (Math.abs(offset) % range); + for (int p = basePort; p <= maxPort; p++) { + if (inUse.add(p)) { + return p; + } + } + throw new IllegalStateException( + "No free ports in range " + basePort + "-" + maxPort); + } + + public void release(int port) { + inUse.remove(port); } } diff --git a/src/main/java/io/github/hectorvent/floci/services/lambda/launcher/ContainerLauncher.java b/src/main/java/io/github/hectorvent/floci/services/lambda/launcher/ContainerLauncher.java index 220f5b5479..16c49b2864 100644 --- a/src/main/java/io/github/hectorvent/floci/services/lambda/launcher/ContainerLauncher.java +++ b/src/main/java/io/github/hectorvent/floci/services/lambda/launcher/ContainerLauncher.java @@ -286,6 +286,8 @@ public void stop(ContainerHandle handle) { } catch (ExecutionException | TimeoutException e) { LOG.warnv(e, "RuntimeApiServer did not close cleanly for container {0}", handle.getContainerId()); + } finally { + runtimeApiServerFactory.release(handle.getRuntimeApiServer()); } lifecycleManager.stopAndRemove(handle.getContainerId(), handle.getLogStream()); } diff --git a/src/main/java/io/github/hectorvent/floci/services/lambda/runtime/RuntimeApiServerFactory.java b/src/main/java/io/github/hectorvent/floci/services/lambda/runtime/RuntimeApiServerFactory.java index dd785153e5..b18862e884 100644 --- a/src/main/java/io/github/hectorvent/floci/services/lambda/runtime/RuntimeApiServerFactory.java +++ b/src/main/java/io/github/hectorvent/floci/services/lambda/runtime/RuntimeApiServerFactory.java @@ -29,16 +29,22 @@ public RuntimeApiServerFactory(Vertx vertx, PortAllocator portAllocator) { public RuntimeApiServer create() { int port = portAllocator.allocate(); - RuntimeApiServer server = new RuntimeApiServer(vertx, port); try { + RuntimeApiServer server = new RuntimeApiServer(vertx, port); server.start().get(10, TimeUnit.SECONDS); LOG.debugv("Created RuntimeApiServer on port {0}", port); + return server; } catch (InterruptedException e) { + portAllocator.release(port); Thread.currentThread().interrupt(); throw new RuntimeException("Interrupted while starting RuntimeApiServer", e); } catch (ExecutionException | TimeoutException e) { + portAllocator.release(port); throw new RuntimeException("Failed to start RuntimeApiServer on port " + port, e); } - return server; + } + + public void release(RuntimeApiServer server) { + portAllocator.release(server.getPort()); } } diff --git a/src/test/java/io/github/hectorvent/floci/core/common/port/PortAllocatorTest.java b/src/test/java/io/github/hectorvent/floci/core/common/port/PortAllocatorTest.java index b90be4bc1f..07685f91ed 100644 --- a/src/test/java/io/github/hectorvent/floci/core/common/port/PortAllocatorTest.java +++ b/src/test/java/io/github/hectorvent/floci/core/common/port/PortAllocatorTest.java @@ -2,6 +2,7 @@ import org.junit.jupiter.api.Test; +import java.util.HashSet; import java.util.Set; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.CountDownLatch; @@ -39,4 +40,25 @@ void concurrentAllocationsAreUnique() throws InterruptedException { executor.shutdown(); assertEquals(threads, ports.size(), "All allocated ports must be unique"); } + + @Test + void allocateNeverReturnsPortAlreadyHandedOut() { + PortAllocator allocator = new PortAllocator(9200, 9209); + Set handed = new HashSet<>(); + for (int i = 0; i < 10; i++) { + assertTrue(handed.add(allocator.allocate())); + } + assertThrows(IllegalStateException.class, allocator::allocate); + } + + @Test + void releasedPortBecomesAvailableAgain() { + PortAllocator allocator = new PortAllocator(9200, 9201); + int first = allocator.allocate(); + allocator.allocate(); + assertThrows(IllegalStateException.class, allocator::allocate); + + allocator.release(first); + assertEquals(first, allocator.allocate()); + } }