✨ feat(database): add health check support for DatabaseApi and registry management

- introduce DatabaseHealthRegistry for managing health checks of DatabaseApi instances
- enhance DatabaseApi constructor to register health check name
- implement DatabaseMicroserviceHealthContributor for microservice health checks
- update shutdown method to unregister from health registry

Author: twisti-dev

build.gradle.kts

@@ -16,6 +16,8 @@ dependencies {
     implementation(libs.bundles.exposed)
     implementation(libs.r2dbc.pool)
     implementation(libs.bundles.databaseDriver)
+
+    compileOnly(libs.surf.microservice)
 }
 
 configurations.runtimeClasspath {

gradle.properties

@@ -3,4 +3,4 @@ kotlin.stdlib.default.dependency=false
 org.gradle.parallel=true
 #org.gradle.caching=true
 #org.gradle.configureondemand=true
-version=2.3.1
+version=2.3.2

gradle/libs.versions.toml

@@ -3,6 +3,7 @@ exposed = "1.3.0"
 r2dbc-mariadb = "1.4.1"
 r2dbc-postgresql = "1.1.1.RELEASE"
 r2dbc-pool = "1.0.2.RELEASE"
+surf-microservice = "2.+"
 
 [libraries]
 exposed-core = { module = "org.jetbrains.exposed:exposed-core", version.ref = "exposed" }
@@ -13,7 +14,7 @@ exposed-migration-r2dbc = { module = "org.jetbrains.exposed:exposed-migration-r2
 r2dbc-mariadb = { module = "org.mariadb:r2dbc-mariadb", version.ref = "r2dbc-mariadb" }
 r2dbc-postgresql = { module = "org.postgresql:r2dbc-postgresql", version.ref = "r2dbc-postgresql" }
 r2dbc-pool = { module = "io.r2dbc:r2dbc-pool", version.ref = "r2dbc-pool" }
-
+surf-microservice = { module = "dev.slne.surf.microservice:surf-microservice-api-microservice", version.ref = "surf-microservice" }
 
 [bundles]
 exposed = ["exposed-core", "exposed-r2dbc", "exposed-java-time", "exposed-json", "exposed-migration-r2dbc"]

gradle/wrapper/gradle-wrapper.properties

@@ -1,6 +1,6 @@
 distributionBase=GRADLE_USER_HOME
 distributionPath=wrapper/dists
-distributionUrl=https\://services.gradle.org/distributions/gradle-9.6.0-bin.zip
+distributionUrl=https\://services.gradle.org/distributions/gradle-9.6.1-bin.zip
 networkTimeout=10000
 retries=0
 retryBackOffMs=500

gradlew

@@ -3,6 +3,7 @@ package dev.slne.surf.database
 import dev.slne.surf.api.core.util.getCallerClass
 import dev.slne.surf.database.config.DatabaseConfig
 import dev.slne.surf.database.config.DatabaseType
+import dev.slne.surf.database.health.DatabaseHealthRegistry
 import dev.slne.surf.database.logger.ComponentSqlLogger
 import io.r2dbc.pool.ConnectionPool
 import io.r2dbc.pool.ConnectionPoolConfiguration
@@ -23,15 +24,28 @@ import org.slf4j.event.Level
 import reactor.netty.resources.LoopResources
 import java.nio.file.Path
 import java.time.Duration.ofMillis
+import java.util.concurrent.atomic.AtomicLong
 
 /**
  * Small wrapper around an [R2dbcDatabase] instance.
  *
  * Use [create(Path, String, R2dbcDatabaseConfig.Builder.() -> Unit)] in production to build the connection pool from the
  * on-disk [DatabaseConfig]. A lower-level overload exists mainly for tests.
  */
-class DatabaseApi internal constructor(val database: R2dbcDatabase) {
+class DatabaseApi internal constructor(
+    val database: R2dbcDatabase,
+    healthCheckName: String
+) {
+    init {
+        DatabaseHealthRegistry.register(
+            databaseApi = this,
+            name = healthCheckName
+        )
+    }
+
     companion object {
+        private val manualDatabaseSequence = AtomicLong()
+
         /**
          * Creates a [DatabaseApi] using the [DatabaseConfig] located in/relative to [pluginPath].
          *
@@ -42,7 +56,6 @@ class DatabaseApi internal constructor(val database: R2dbcDatabase) {
          * @param poolName Optional pool name (defaults to a generated name based on the caller).
          * @param configCustomizer Optional customization hook for Exposed's [R2dbcDatabaseConfig].
          */
-        @OptIn(TestOnlyDatabaseApi::class)
         fun create(
             pluginPath: Path,
             fileName: String = "database.yml",
@@ -121,11 +134,12 @@ class DatabaseApi internal constructor(val database: R2dbcDatabase) {
                 DatabaseType.POSTGRESQL -> PostgreSQLDialect()
             }
 
-            return create(
+            return createDatabase(
                 connectionFactory = pool,
                 dialect = dialect,
                 logger = logger,
                 logLevel = logLevel,
+                healthCheckName = poolName,
                 configCustomizer = configCustomizer
             )
         }
@@ -151,21 +165,50 @@ class DatabaseApi internal constructor(val database: R2dbcDatabase) {
             logLevel: Level = Level.DEBUG,
             configCustomizer: R2dbcDatabaseConfig.Builder.() -> Unit = {}
         ): DatabaseApi {
-            val database = R2dbcDatabase.connect(connectionFactory, R2dbcDatabaseConfig {
-                explicitDialect = dialect
-                sqlLogger = ComponentSqlLogger(logger, logLevel)
-                defaultR2dbcIsolationLevel = IsolationLevel.READ_UNCOMMITTED
-                configCustomizer()
-            })
-
-            return DatabaseApi(database)
+            return createDatabase(
+                connectionFactory = connectionFactory,
+                dialect = dialect,
+                logger = logger,
+                logLevel = logLevel,
+                healthCheckName = nextManualHealthCheckName(),
+                configCustomizer = configCustomizer
+            )
+        }
+
+        private fun createDatabase(
+            connectionFactory: ConnectionFactory,
+            dialect: DatabaseDialect,
+            logger: ComponentLogger,
+            logLevel: Level,
+            healthCheckName: String,
+            configCustomizer: R2dbcDatabaseConfig.Builder.() -> Unit
+        ): DatabaseApi {
+            val database = R2dbcDatabase.connect(
+                connectionFactory,
+                R2dbcDatabaseConfig {
+                    explicitDialect = dialect
+                    sqlLogger = ComponentSqlLogger(logger, logLevel)
+                    defaultR2dbcIsolationLevel = IsolationLevel.READ_UNCOMMITTED
+
+                    configCustomizer()
+                }
+            )
+
+            return DatabaseApi(
+                database = database,
+                healthCheckName = healthCheckName
+            )
         }
 
         private fun generatePoolName(): String {
             val caller = getCallerClass(1)
             val callerName = caller?.simpleName ?: "unknown"
             return "j2bdc-pool-$callerName"
         }
+
+        private fun nextManualHealthCheckName(): String {
+            return "database-${manualDatabaseSequence.incrementAndGet()}"
+        }
     }
 
     /**
@@ -174,7 +217,11 @@ class DatabaseApi internal constructor(val database: R2dbcDatabase) {
      * Call this during plugin/application shutdown to ensure connections are released.
      */
     fun shutdown() {
-        TransactionManager.closeAndUnregister(database)
+        try {
+            TransactionManager.closeAndUnregister(database)
+        } finally {
+            DatabaseHealthRegistry.unregister(this)
+        }
     }
 }
 

gradlew.bat

@@ -0,0 +1,32 @@
+package dev.slne.surf.database.health
+
+import dev.slne.surf.database.DatabaseApi
+import java.util.concurrent.ConcurrentHashMap
+
+internal data class RegisteredDatabase(
+    val name: String,
+    val api: DatabaseApi
+)
+
+internal object DatabaseHealthRegistry {
+    private val databases = ConcurrentHashMap<DatabaseApi, String>()
+
+    fun register(databaseApi: DatabaseApi, name: String) {
+        databases[databaseApi] = name
+    }
+
+    fun unregister(databaseApi: DatabaseApi) {
+        databases.remove(databaseApi)
+    }
+
+    fun snapshot(): List<RegisteredDatabase> {
+        return databases.entries
+            .map { (api, name) ->
+                RegisteredDatabase(
+                    name = name,
+                    api = api
+                )
+            }
+            .sortedBy(RegisteredDatabase::name)
+    }
+}

src/main/kotlin/dev/slne/surf/database/DatabaseApi.kt

@@ -0,0 +1,82 @@
+package dev.slne.surf.database.health
+
+import com.google.auto.service.AutoService
+import dev.slne.surf.microservice.api.microservice.health.MicroserviceHealthCheckResult
+import dev.slne.surf.microservice.api.microservice.health.MicroserviceHealthContributor
+import dev.slne.surf.microservice.api.microservice.health.MicroserviceHealthStatus
+import kotlinx.coroutines.async
+import kotlinx.coroutines.awaitAll
+import kotlinx.coroutines.supervisorScope
+import kotlinx.coroutines.withTimeoutOrNull
+import org.jetbrains.exposed.v1.r2dbc.transactions.suspendTransaction
+import kotlin.coroutines.cancellation.CancellationException
+import kotlin.time.Duration.Companion.seconds
+
+private val DATABASE_HEALTH_TIMEOUT = 2.5.seconds
+
+@AutoService(MicroserviceHealthContributor::class)
+internal class DatabaseMicroserviceHealthContributor : MicroserviceHealthContributor {
+    override val name: String = "database"
+
+    override suspend fun check(): List<MicroserviceHealthCheckResult> {
+        val databases = DatabaseHealthRegistry.snapshot()
+
+        return supervisorScope {
+            databases
+                .map { database ->
+                    async {
+                        checkDatabase(database)
+                    }
+                }
+                .awaitAll()
+        }
+    }
+
+    private suspend fun checkDatabase(
+        database: RegisteredDatabase
+    ): MicroserviceHealthCheckResult = try {
+        val completed = withTimeoutOrNull(DATABASE_HEALTH_TIMEOUT) {
+            suspendTransaction(db = database.api.database) {
+                exec("SELECT 1")
+            }
+
+            true
+        } ?: false
+
+        if (completed) {
+            healthyResult(database.name)
+        } else {
+            unhealthyResult(
+                instance = database.name,
+                message = "Database health check timed out after $DATABASE_HEALTH_TIMEOUT."
+            )
+        }
+    } catch (exception: CancellationException) {
+        throw exception
+    } catch (exception: Throwable) {
+        unhealthyResult(
+            instance = database.name,
+            message = exception.message
+                ?.takeIf { it.isNotBlank() }
+                ?: exception.javaClass.name
+        )
+    }
+
+    private fun healthyResult(
+        instance: String
+    ) = MicroserviceHealthCheckResult(
+        component = name,
+        instance = instance,
+        status = MicroserviceHealthStatus.HEALTHY
+    )
+
+    private fun unhealthyResult(
+        instance: String,
+        message: String
+    ) = MicroserviceHealthCheckResult(
+        component = name,
+        instance = instance,
+        status = MicroserviceHealthStatus.UNHEALTHY,
+        message = message
+    )
+}