refactor!: Share single Spanner emulator process (#283)

The Spanner emulator documentation recommends this. It should hopefully reduce test flakiness.

BREAKING CHANGE: The `SpannerEmulatorDatabaseRule` constructor has a new required parameter of type `SpannerDatabaseAdmin`.

src/main/kotlin/org/wfanet/measurement/gcloud/spanner/testing/SpannerEmulatorDatabaseRule.kt

@@ -14,67 +14,40 @@
 
 package org.wfanet.measurement.gcloud.spanner.testing
 
-import com.google.cloud.spanner.DatabaseId
-import com.google.cloud.spanner.Spanner
-import com.google.cloud.spanner.connection.SpannerPool
 import java.nio.file.Path
-import java.sql.DriverManager
-import kotlinx.coroutines.runBlocking
+import java.util.concurrent.atomic.AtomicInteger
 import org.junit.rules.TestRule
 import org.junit.runner.Description
 import org.junit.runners.model.Statement
-import org.wfanet.measurement.common.db.liquibase.Liquibase
 import org.wfanet.measurement.gcloud.spanner.AsyncDatabaseClient
-import org.wfanet.measurement.gcloud.spanner.buildSpanner
-import org.wfanet.measurement.gcloud.spanner.getAsyncDatabaseClient
 
 /**
  * JUnit rule exposing a temporary Google Cloud Spanner database via Spanner Emulator.
  *
  * @param changelogPath [Path] to a Liquibase changelog.
  */
 class SpannerEmulatorDatabaseRule(
+  private val emulatorDatabaseAdmin: SpannerDatabaseAdmin,
   private val changelogPath: Path,
-  private val databaseName: String = "test-db",
+  private val databaseId: String = "test-db-" + dbCounter.incrementAndGet(),
 ) : TestRule {
   lateinit var databaseClient: AsyncDatabaseClient
     private set
 
   override fun apply(base: Statement, description: Description): Statement {
     return object : Statement() {
       override fun evaluate() {
-        check(!::databaseClient.isInitialized)
-
-        SpannerEmulator().use { emulator ->
-          val emulatorHost = runBlocking { emulator.start() }
-          try {
-            createDatabase(emulatorHost).use { spanner ->
-              databaseClient =
-                spanner.getAsyncDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, databaseName))
-              base.evaluate()
-            }
-          } finally {
-            // Make sure these Spanner instances from JDBC are closed before the emulator is shut
-            // down, otherwise it will block JVM shutdown.
-            SpannerPool.closeSpannerPool()
-          }
+        databaseClient = emulatorDatabaseAdmin.createDatabase(changelogPath, databaseId)
+        try {
+          base.evaluate()
+        } finally {
+          emulatorDatabaseAdmin.deleteDatabase(databaseId)
         }
       }
     }
   }
 
-  private fun createDatabase(emulatorHost: String): Spanner {
-    val connectionString =
-      SpannerEmulator.buildJdbcConnectionString(emulatorHost, PROJECT, INSTANCE, databaseName)
-    DriverManager.getConnection(connectionString).use { connection ->
-      Liquibase.update(connection, changelogPath)
-    }
-
-    return buildSpanner(PROJECT, emulatorHost)
-  }
-
   companion object {
-    private const val PROJECT = "test-project"
-    private const val INSTANCE = "test-instance"
+    private val dbCounter = AtomicInteger(0)
   }
 }

src/main/kotlin/org/wfanet/measurement/gcloud/spanner/testing/SpannerEmulatorRule.kt

@@ -0,0 +1,99 @@
+/*
+ * Copyright 2024 The Cross-Media Measurement Authors
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.wfanet.measurement.gcloud.spanner.testing
+
+import com.google.cloud.spanner.DatabaseId
+import com.google.cloud.spanner.Spanner
+import com.google.cloud.spanner.connection.SpannerPool
+import java.nio.file.Path
+import java.sql.DriverManager
+import kotlinx.coroutines.runBlocking
+import org.junit.rules.TestRule
+import org.junit.runner.Description
+import org.junit.runners.model.Statement
+import org.wfanet.measurement.common.db.liquibase.Liquibase
+import org.wfanet.measurement.gcloud.spanner.AsyncDatabaseClient
+import org.wfanet.measurement.gcloud.spanner.buildSpanner
+import org.wfanet.measurement.gcloud.spanner.getAsyncDatabaseClient
+
+/** Administration of databases within a Spanner instance. */
+interface SpannerDatabaseAdmin {
+  /** Creates a database. */
+  fun createDatabase(changelogPath: Path, databaseId: String): AsyncDatabaseClient
+
+  /** Deletes a database. */
+  fun deleteDatabase(databaseId: String)
+}
+
+/**
+ * [TestRule] which manages a [SpannerEmulator] resource.
+ *
+ * This is intended to be used as a [org.junit.ClassRule].
+ */
+class SpannerEmulatorRule : TestRule, SpannerDatabaseAdmin {
+  private lateinit var emulatorHost: String
+  private lateinit var spanner: Spanner
+
+  override fun apply(base: Statement, description: Description): Statement {
+    return object : Statement() {
+      override fun evaluate() {
+        SpannerEmulator().use { emulator ->
+          try {
+            emulatorHost = runBlocking { emulator.start() }
+            base.evaluate()
+          } finally {
+            if (::spanner.isInitialized) {
+              spanner.close()
+            }
+            // Make sure these Spanner instances from JDBC are closed before the emulator is shut
+            // down, otherwise it will block JVM shutdown.
+            SpannerPool.closeSpannerPool()
+          }
+        }
+      }
+    }
+  }
+
+  override fun createDatabase(changelogPath: Path, databaseId: String): AsyncDatabaseClient {
+    check(::emulatorHost.isInitialized) {
+      "Spanner emulator has not been started. " +
+        "Ensure that SpannerEmulatorRule has been registered as a ClassRule."
+    }
+
+    val connectionString =
+      SpannerEmulator.buildJdbcConnectionString(emulatorHost, PROJECT, INSTANCE, databaseId)
+    DriverManager.getConnection(connectionString).use { connection ->
+      Liquibase.update(connection, changelogPath)
+    }
+
+    // Spanner must be initialized only after the first database has been created.
+    if (!::spanner.isInitialized) {
+      spanner = buildSpanner(PROJECT, emulatorHost)
+    }
+
+    return spanner.getAsyncDatabaseClient(DatabaseId.of(PROJECT, INSTANCE, databaseId))
+  }
+
+  override fun deleteDatabase(databaseId: String) {
+    spanner.databaseAdminClient.dropDatabase(INSTANCE, databaseId)
+  }
+
+  companion object {
+    private const val PROJECT = "test-project"
+    private const val INSTANCE = "test-instance"
+  }
+}

src/main/kotlin/org/wfanet/measurement/gcloud/spanner/testing/UsingSpannerEmulator.kt

@@ -15,6 +15,7 @@
 package org.wfanet.measurement.gcloud.spanner.testing
 
 import java.nio.file.Path
+import org.junit.ClassRule
 import org.junit.Rule
 import org.wfanet.measurement.gcloud.spanner.AsyncDatabaseClient
 
@@ -40,8 +41,13 @@ import org.wfanet.measurement.gcloud.spanner.AsyncDatabaseClient
  * ```
  */
 abstract class UsingSpannerEmulator(changeLogResourcePath: Path) {
-  @get:Rule val spannerDatabase = SpannerEmulatorDatabaseRule(changeLogResourcePath)
+  @get:Rule
+  val spannerDatabase = SpannerEmulatorDatabaseRule(spannerEmulator, changeLogResourcePath)
 
   val databaseClient: AsyncDatabaseClient
     get() = spannerDatabase.databaseClient
+
+  companion object {
+    @get:ClassRule @JvmStatic val spannerEmulator = SpannerEmulatorRule()
+  }
 }

src/main/kotlin/org/wfanet/measurement/gcloud/spanner/testing/macros.bzl

@@ -16,9 +16,15 @@
 
 load("@wfa_rules_kotlin_jvm//kotlin:defs.bzl", "kt_jvm_test")
 
-def spanner_emulator_test(name, data = [], **kwargs):
+def spanner_emulator_test(name, data = None, tags = None, **kwargs):
+    data = data or []
+    tags = tags or []
     kt_jvm_test(
         name = name,
         data = data + ["@cloud_spanner_emulator//:emulator"],
+        tags = tags + [
+            # Only one Spanner emulator process should be running at a time.
+            "exclusive",
+        ],
         **kwargs
     )

src/test/kotlin/org/wfanet/measurement/gcloud/spanner/AsyncDatabaseClientTest.kt

@@ -21,24 +21,21 @@ import com.google.common.truth.Truth.assertThat
 import java.nio.file.Path
 import kotlinx.coroutines.flow.toList
 import kotlinx.coroutines.runBlocking
-import org.junit.Before
+import org.junit.ClassRule
 import org.junit.Rule
 import org.junit.Test
 import org.junit.runner.RunWith
 import org.junit.runners.JUnit4
 import org.wfanet.measurement.common.getJarResourcePath
 import org.wfanet.measurement.gcloud.spanner.testing.SpannerEmulatorDatabaseRule
+import org.wfanet.measurement.gcloud.spanner.testing.SpannerEmulatorRule
 
 @RunWith(JUnit4::class)
 class AsyncDatabaseClientTest {
-  @JvmField @Rule val spannerEmulatorDb = SpannerEmulatorDatabaseRule(CHANGELOG_PATH)
+  @get:Rule val database = SpannerEmulatorDatabaseRule(spannerEmulator, CHANGELOG_PATH)
 
-  private lateinit var databaseClient: AsyncDatabaseClient
-
-  @Before
-  fun initDatabaseClient() {
-    databaseClient = spannerEmulatorDb.databaseClient
-  }
+  private val databaseClient: AsyncDatabaseClient
+    get() = database.databaseClient
 
   @Test
   fun `executes simple query`() {
@@ -55,5 +52,7 @@ class AsyncDatabaseClientTest {
       requireNotNull(this::class.java.classLoader.getJarResourcePath(CHANGELOG_RESOURCE_NAME)) {
         "Resource $CHANGELOG_RESOURCE_NAME not found"
       }
+
+    @get:ClassRule @JvmStatic val spannerEmulator = SpannerEmulatorRule()
   }
 }

src/test/kotlin/org/wfanet/measurement/gcloud/spanner/BUILD.bazel

@@ -4,6 +4,10 @@ kt_jvm_test(
     name = "AsyncDatabaseClientTest",
     srcs = ["AsyncDatabaseClientTest.kt"],
     resources = ["//src/test/resources/db/spanner"],
+    tags = [
+        # There should only be one Spanner emulator process running.
+        "exclusive",
+    ],
     test_class = "org.wfanet.measurement.gcloud.spanner.AsyncDatabaseClientTest",
     deps = [
         "//imports/java/com/google/cloud/spanner",