Try to close reused sockets in TestDynamoDbService

cashapp · Jun 3, 2024 · a0284f2 · a0284f2
1 parent 4b6b047
commit a0284f2
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 20 deletions.
diff --git a/...testing-internal/src/main/kotlin/app/cash/tempest/testing/internal/TestDynamoDbService.kt b/...testing-internal/src/main/kotlin/app/cash/tempest/testing/internal/TestDynamoDbService.kt
@@ -4,6 +4,7 @@ import app.cash.tempest.testing.TestDynamoDbClient
 import app.cash.tempest.testing.TestDynamoDbServer
 import app.cash.tempest.testing.TestTable
 import com.google.common.util.concurrent.AbstractIdleService
+import java.net.ServerSocket
 import java.util.concurrent.ConcurrentHashMap
 
 /**
@@ -45,24 +46,27 @@ class TestDynamoDbService private constructor(
   }
 
   companion object {
-    private val defaultPorts = ConcurrentHashMap<String, Int>()
-    private fun defaultPort(key: String): PortHolder {
-      var releasePort: () -> Unit = {}
+    private val allocatedSockets = ConcurrentHashMap<String, ServerSocket>()
+    private fun defaultPortHolder(key: String): PortHolder {
       // Only pick random port once to share one test server with multiple tests.
-      val port = defaultPorts.getOrPut(key) {
-        val socket = allocateRandomPort()
-        releasePort = { socket.close() }
-        socket.localPort
+      val socket = allocatedSockets.getOrPut(key) { allocateRandomPort() }
+      return PortHolder(socket.localPort) {
+        if (!socket.isClosed) {
+          socket.close()
+        }
       }
-      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 portHolder = port?.let { PortHolder(it) } ?: defaultPort(serverFactory.toString())
+    fun create(
+      serverFactory: TestDynamoDbServer.Factory<*>,
+      tables: List<TestTable>,
+      port: Int? = null
+    ): TestDynamoDbService {
+      val portHolder = port?.let { PortHolder(it) } ?: defaultPortHolder(serverFactory.toString())
       return TestDynamoDbService(
         DefaultTestDynamoDbClient(tables, portHolder.value),
         serverFactory.create(portHolder.value, portHolder.releasePort)

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
@@ -32,6 +32,10 @@ import java.net.InetSocketAddress
 import java.net.ServerSocket
 import java.net.Socket
 
+fun pickRandomPort(): Int {
+  ServerSocket(0).use { socket -> return socket.localPort }
+}
+
 fun allocateRandomPort(): ServerSocket {
   val socket = ServerSocket(0)
   Runtime.getRuntime().addShutdownHook(

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
@@ -4,6 +4,7 @@ import app.cash.tempest2.testing.TestDynamoDbClient
 import app.cash.tempest2.testing.TestDynamoDbServer
 import app.cash.tempest2.testing.TestTable
 import com.google.common.util.concurrent.AbstractIdleService
+import java.net.ServerSocket
 import java.util.concurrent.ConcurrentHashMap
 
 /**
@@ -45,24 +46,27 @@ class TestDynamoDbService(
   }
 
   companion object {
-    private val defaultPorts = ConcurrentHashMap<String, Int>()
-    private fun defaultPort(key: String): PortHolder {
-      var releasePort: () -> Unit = {}
+    private val allocatedSockets = ConcurrentHashMap<String, ServerSocket>()
+    private fun defaultPortHolder(key: String): PortHolder {
       // Only pick random port once to share one test server with multiple tests.
-      val port = defaultPorts.getOrPut(key) {
-        val socket = allocateRandomPort()
-        releasePort = { socket.close() }
-        socket.localPort
+      val socket = allocatedSockets.getOrPut(key) { allocateRandomPort() }
+      return PortHolder(socket.localPort) {
+        if (!socket.isClosed) {
+          socket.close()
+        }
       }
-      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 portHolder = port?.let { PortHolder(it) } ?: defaultPort(serverFactory.toString())
+    fun create(
+      serverFactory: TestDynamoDbServer.Factory<*>,
+      tables: List<TestTable>,
+      port: Int? = null
+    ): TestDynamoDbService {
+      val portHolder = port?.let { PortHolder(it) } ?: defaultPortHolder(serverFactory.toString())
       return TestDynamoDbService(
         DefaultTestDynamoDbClient(tables, portHolder.value),
         serverFactory.create(portHolder.value, portHolder.releasePort)

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,6 +33,10 @@ 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(