Tempest2 - Hold onto allocated ports to avoid port conflicts in concu…

…rrent tests (cashapp#184)
mmollaverdi · May 31, 2024 · 4b6b047 · 4b6b047
1 parent 38c0913
commit 4b6b047
Show file tree

Hide file tree

Showing 6 changed files with 30 additions and 13 deletions.
diff --git a/tempest-testing-internal/src/main/kotlin/app/cash/tempest/testing/internal/TestUtils.kt b/tempest-testing-internal/src/main/kotlin/app/cash/tempest/testing/internal/TestUtils.kt
@@ -33,7 +33,7 @@ import java.net.ServerSocket
 import java.net.Socket
 
 fun allocateRandomPort(): ServerSocket {
-  val socket = ServerSocket(0) //use { socket -> return socket.localPort }
+  val socket = ServerSocket(0)
   Runtime.getRuntime().addShutdownHook(
     Thread { socket.close() }
   )

diff --git a/tempest2-testing-docker/src/main/kotlin/app/cash/tempest2/testing/DockerDynamoDbServer.kt b/tempest2-testing-docker/src/main/kotlin/app/cash/tempest2/testing/DockerDynamoDbServer.kt
@@ -24,12 +24,14 @@ import software.amazon.awssdk.services.dynamodb.model.DeleteTableRequest
 import software.amazon.awssdk.services.dynamodb.model.DynamoDbException
 
 class DockerDynamoDbServer private constructor(
-  override val port: Int
+  override val port: Int,
+  private val onBeforeStartup: () -> Unit
 ) : AbstractIdleService(), TestDynamoDbServer {
 
   override val id = "tempest2-docker-dynamodb-local-$port"
 
   override fun startUp() {
+    onBeforeStartup()
     composer.start()
 
     // Temporary client to block until the container is running
@@ -67,6 +69,6 @@ class DockerDynamoDbServer private constructor(
   )
 
   object Factory : TestDynamoDbServer.Factory<DockerDynamoDbServer> {
-    override fun create(port: Int) = DockerDynamoDbServer(port)
+    override fun create(port: Int, onBeforeStartup: () -> Unit) = DockerDynamoDbServer(port, onBeforeStartup)
   }
 }
diff --git a/...esting-internal/src/main/kotlin/app/cash/tempest2/testing/internal/TestDynamoDbService.kt b/...esting-internal/src/main/kotlin/app/cash/tempest2/testing/internal/TestDynamoDbService.kt
@@ -46,21 +46,29 @@ class TestDynamoDbService(
 
   companion object {
     private val defaultPorts = ConcurrentHashMap<String, Int>()
-    private fun defaultPort(key: String): Int {
+    private fun defaultPort(key: String): PortHolder {
+      var releasePort: () -> Unit = {}
       // Only pick random port once to share one test server with multiple tests.
-      return defaultPorts.getOrPut(key, ::pickRandomPort)
+      val port = defaultPorts.getOrPut(key) {
+        val socket = allocateRandomPort()
+        releasePort = { socket.close() }
+        socket.localPort
+      }
+      return PortHolder(port, releasePort)
     }
 
     private val runningServers = ConcurrentHashMap.newKeySet<String>()
     private val log = getLogger<TestDynamoDbService>()
 
     @JvmStatic
     fun create(serverFactory: TestDynamoDbServer.Factory<*>, tables: List<TestTable>, port: Int? = null): TestDynamoDbService {
-      val port = port ?: defaultPort(serverFactory.toString())
+      val portHolder = port?.let { PortHolder(it) } ?: defaultPort(serverFactory.toString())
       return TestDynamoDbService(
-        DefaultTestDynamoDbClient(tables, port),
-        serverFactory.create(port)
+        DefaultTestDynamoDbClient(tables, portHolder.value),
+        serverFactory.create(portHolder.value, portHolder.releasePort)
       )
     }
   }
+
+  private data class PortHolder(val value: Int, val releasePort: () -> Unit = {})
 }
diff --git a/tempest2-testing-internal/src/main/kotlin/app/cash/tempest2/testing/internal/TestUtils.kt b/tempest2-testing-internal/src/main/kotlin/app/cash/tempest2/testing/internal/TestUtils.kt
@@ -33,8 +33,12 @@ import java.net.ServerSocket
 import java.net.Socket
 import java.net.URI
 
-fun pickRandomPort(): Int {
-  ServerSocket(0).use { socket -> return socket.localPort }
+fun allocateRandomPort(): ServerSocket {
+  val socket = ServerSocket(0)
+  Runtime.getRuntime().addShutdownHook(
+    Thread { socket.close() }
+  )
+  return socket
 }
 
 private val AWS_CREDENTIALS_PROVIDER = StaticCredentialsProvider.create(

diff --git a/tempest2-testing-jvm/src/main/kotlin/app/cash/tempest2/testing/JvmDynamoDbServer.kt b/tempest2-testing-jvm/src/main/kotlin/app/cash/tempest2/testing/JvmDynamoDbServer.kt
@@ -22,7 +22,8 @@ import com.google.common.util.concurrent.AbstractIdleService
 import java.io.File
 
 class JvmDynamoDbServer private constructor(
-  override val port: Int
+  override val port: Int,
+  private val onBeforeStartup: () -> Unit
 ) : AbstractIdleService(), TestDynamoDbServer {
 
   override val id = "tempest2-jvm-dynamodb-local-$port"
@@ -33,6 +34,7 @@ class JvmDynamoDbServer private constructor(
     val libraryFile = libsqlite4javaNativeLibrary()
     System.setProperty("sqlite4java.library.path", libraryFile.parent)
 
+    onBeforeStartup()
     server = ServerRunner.createServerFromCommandLineArgs(
       arrayOf("-inMemory", "-port", port.toString())
     )
@@ -87,6 +89,6 @@ class JvmDynamoDbServer private constructor(
   }
 
   object Factory : TestDynamoDbServer.Factory<JvmDynamoDbServer> {
-    override fun create(port: Int) = JvmDynamoDbServer(port)
+    override fun create(port: Int, onBeforeStartup: () -> Unit) = JvmDynamoDbServer(port, onBeforeStartup)
   }
 }
diff --git a/tempest2-testing/src/main/kotlin/app/cash/tempest2/testing/TestDynamoDbServer.kt b/tempest2-testing/src/main/kotlin/app/cash/tempest2/testing/TestDynamoDbServer.kt
@@ -26,6 +26,7 @@ interface TestDynamoDbServer : Service {
   val port: Int
 
   interface Factory<T : TestDynamoDbServer> {
-    fun create(port: Int): T
+    fun create(port: Int): T = create(port) {}
+    fun create(port: Int, onBeforeStartup: () -> Unit): T
   }
 }