Optimistic & Pessimistic locking in mongoDB with Kotlin

Need of Locking:

The locking mechanism is used for concurrency control

When concurrent operations are happening on database there might be the case, when two requests/tasks read the data at same time & trying to update the data which may result in data inconsistency.

Consider example of money transactions. Although real-life money transactions does not follow below approach, but let’s consider below example for making it simple.

Suppose the account balance is $100 and two concurrent transactions are happening on same account. Transaction A is to add $20 in account & Transaction B is to deduct $50 from account

As both transactions were concurrent they read the balance from account at same time i.e $100 and both will try to update the balance , So transaction A will update it to $120 & transaction B will update it to $50 which will result in wrong Account balance

To solve this problem we can use Locking mechanism

There are two types of locking mechanism

1. Optimistic Locking
2. Pessimistic Locking

1. Optimistic Locking

In this mechanism, we allow transactions to happen concurrently but while updating the data we check if data is already updated , if not updated we complete the write operation , otherwise we through the exception

So in our example, Both transaction A & B will read the balance as $100, if transaction A is completed, the updated balance will be $120. Now, when transaction B tries to update the data, first it checks if the balance is updated or not, if it is already updated transaction B does not update balance & we get the exception

Optimistic locking can be achieved using version numbers, timestamps etc.

Let’s see the implementation of Optimistic locking in mongo with kotlin using version numbers.

We can use @Version annotation provided by spring framework in the Mongo Document to implement Optimistic concurrency. While updating the data in db it will check if version is already updated, if yes it will throw OptimisticLockingFailureException, otherwise it will update the data.

Model :

@Document(collection = "accounts")
data class Account(
    @Id
    val id: String? = null,

    var balance: Double,

    @Version
    val version: Long? = null
)

Repository :

interface AccountRepository : MongoRepository<Account, String>

Service :

@Service
class AccountService(@Autowired  val accountRepository: AccountRepository) {

    fun getAccountById(id: String): Account?{
        return accountRepository.findById(id).orElse(null)
    }

    fun deposit(id: String, amount: Double): Account?{
        val account = accountRepository.findById(id).orElseThrow{ RuntimeException("Account not found") }
        account.balance+=amount

        return try{
            accountRepository.save(account)
        }catch (exception: OptimisticLockingFailureException){
            println("Optimistic lock exception: ${exception.message}")
            null
        }
    }

    fun withdraw(id: String, amount: Double): Account?{
        val account = accountRepository.findById(id).orElseThrow{ RuntimeException("Account not found") }
        account.balance-=amount

        return try{
            accountRepository.save(account)
        }catch (exception: OptimisticLockingFailureException){
            println("Optimistic lock exception: ${exception.message}")
            null
        }
    }
}

Controller :

@RestController
@RequestMapping("/accounts")
class AccountController(@Autowired val accountService: AccountService) {

    @GetMapping("/{id}")
    fun getAccount(@PathVariable id: String): ResponseEntity<Account>{
        val account = accountService.getAccountById(id)?: return ResponseEntity.notFound().build()
        return ResponseEntity.ok(account)
    }

    @PostMapping("/{id}/deposit")
    fun deposit(@PathVariable id: String, @RequestBody transactionRequest: TransactionRequest): ResponseEntity<Account>{
        val account = accountService.deposit(id, transactionRequest.amount)?: return ResponseEntity.status(HttpStatus.CONFLICT).build()
        return ResponseEntity.ok(account)
    }

    @PostMapping("/{id}/withdraw")
    fun withdraw(@PathVariable id: String, @RequestBody transactionRequest: TransactionRequest): ResponseEntity<Account>{
        val account = accountService.withdraw(id, transactionRequest.amount)?: return ResponseEntity.status(HttpStatus.CONFLICT).build()
        return ResponseEntity.ok(account)
    }
}

data class TransactionRequest(val amount: Double)

2. Pessimistic Locking

In this mechanism, when we need to update the document, we first acquire a lock on it, then proceed with the update, and release the lock afterward. This way, the lock is held by only one update request at a time until it is released.

So in our example, If transaction A reads the balance first we acquire lock, Now we will not allow any other transaction to read the balance until the transaction A completes the write operation i.e updates the data

As mongoDB does not support row-level locking mechanisms, pessimistic locking implementation will not be straight forward.

However, we can implement pessimistic locking by using certain techniques like using locked flag or having lock period based on the requirement.

Let’s see the implementation of pessimistic lock using the locked flag.

Model :

@Document(collection = "accounts")
data class Account(
    @Id
    val id: String? = null,
    var balance: Double,
    var locked: Boolean = false
)

Repository :

interface AccountRepository : MongoRepository<Account, String>

Service :

@Service
class AccountService(
    @Autowired val accountRepository: AccountRepository,
    @Autowired val mongoTemplate: MongoTemplate
) {
    fun getAccountById(id: String): Account? {
        return accountRepository.findById(id).orElse(null)
    }

    fun deposit(id: String, amount: Double): Account? {
        val account = lockAccount(id) ?: throw RuntimeException("Account not found or already locked")

        account.balance += amount

        val updatedAccount = accountRepository.save(account)

        unlockAccount(account)

        return updatedAccount
    }

    fun withdraw(id: String, amount: Double): Account? {
        val account = lockAccount(id) ?: throw RuntimeException("Account not found or already locked")

        if (account.balance < amount) {
            unlockAccount(account)
            throw RuntimeException("Insufficient funds")
        }

        account.balance -= amount

        val updatedAccount = accountRepository.save(account)

        unlockAccount(account)

        return updatedAccount
    }

    private fun lockAccount(id: String): Account? {
        val query = Query(Criteria.where("id").`is`(id).and("locked").`is`(false))
        val update = Update().set("locked", true)

        return mongoTemplate.findAndModify(query, update, Account::class.java)
    }

    private fun unlockAccount(account: Account) {
        account.locked = false
        accountRepository.save(account)
    }
}

Controller :

@RestController
@RequestMapping("/accounts")
class AccountController(@Autowired val accountService: AccountService) {
    @GetMapping("/{id}")
    fun getAccount(@PathVariable id: String): ResponseEntity<Account> {
        val account = accountService.getAccountById(id) ?: return ResponseEntity.notFound().build()
        return ResponseEntity.ok(account)
    }

    @PostMapping("/{id}/deposit")
    fun deposit(@PathVariable id: String, @RequestBody transaction: TransactionRequest): ResponseEntity<Account> {
        return try {
            val updatedAccount = accountService.deposit(id, transaction.amount)
            ResponseEntity.ok(updatedAccount)
        } catch (e: RuntimeException) {
            ResponseEntity.status(409).body(null)
        }
    }

    @PostMapping("/{id}/withdraw")
    fun withdraw(@PathVariable id: String, @RequestBody transaction: TransactionRequest): ResponseEntity<Account> {
        return try {
            val updatedAccount = accountService.withdraw(id, transaction.amount)
            ResponseEntity.ok(updatedAccount)
        } catch (e: RuntimeException) {
            ResponseEntity.status(409).body(null)
        }
    }
}

data class TransactionRequest(val amount: Double)

Summary

Optimistic locking is suitable for read-heavy applications, to provide better performance and flexibility whereas Pessimistic locking is suitable for write-heavy applications but at cost of performance & scalability.