Postgres transaction-level advisory lock implementation that uses Spring JDBC

LockManager.kt

import org.slf4j.LoggerFactory
import org.springframework.jdbc.core.JdbcTemplate
import org.springframework.stereotype.Component
import org.springframework.transaction.annotation.Propagation
import org.springframework.transaction.annotation.Transactional
import java.time.Duration

interface LockManager {
    fun <T> tryWithLock(key: Long, timeout: Duration, function: () -> T): T
}

/**
 * LockManager implementation that uses postgres transaction-level advisory locks.
 */
@Component
class PostgresLockManager(private val jdbcTemplate: JdbcTemplate) : LockManager {
    private val log = LoggerFactory.getLogger(LockManager::class.java)

    /**
     * We need a transaction for the locking to actually work, but possible to participate in an ongoing transaction.
     */
    @Transactional(propagation = Propagation.REQUIRED)
    override fun <T> tryWithLock(key: Long, timeout: Duration, function: () -> T): T {
        lock(key, timeout)
        return function()
    }

    private fun lock(key: Long, timeout: Duration) {
        val timeoutMillis = timeout.toMillis()
        var count = 0

        log.debug("Acquiring pg_try_advisory_xact_lock($key)")

        while (!jdbcTemplate.queryForObject("select pg_try_advisory_xact_lock(?)", Boolean::class.java, key)) {
            if (timeoutMillis - (1000 * count++) > 0) {
                log.info("Waiting for 1000 ms to acquire pg_try_advisory_xact_lock($key)")
                Thread.sleep(1000)
            } else {
                throw RuntimeException("Deadlock detected while acquiring pg_try_advisory_xact_lock($key)")
            }
        }
    }
}

/**
 * Dummy implementation for test purposes.
 */
class DummyLockManager : LockManager {
    override fun <T> tryWithLock(key: Long, timeout: Duration, function: () -> T): T = function()
}

/**
 * Example of usage.
 */
@Transactional
class ClassThatNeedsLocking(private val lockManager: LockManager) {

    fun functionThatUsesTheLockManager() {
        lockManager.tryWithLock(1337, Duration.ofMillis(5000)) {
            // do stuff on entity with id=1337 that needs locking.
        }
    }
}