Distributed Locks from scratch
Purpose
The purpose of this document is to implement distributed locking across multiple JVMs without using any framework.
Use Case
The problem with both spring-integration's version of distributed locks or shedlock is that we are reliant on spring.
For an application running an older version of spring or mongo java library, upgrading is generally a huge task & finding the right versions to make your application work is a frustrating process and sometimes there is no right version that works.
Let's write our own version of distributed locks.
Requirement
A scheduled task can be performed only if it has the lock.
Pseudocode
boolean g = grabthelock(lockkey);
if (g) {
// grabbed the lock, ready to perform some work.
performwork();
releasethelock(lockkey);
} else {
// could not grab the lock, cheers.
}
That's all we need!
Things to note:
- grabthelock(lockkey) needs to be an atomic operation - findAndModify.
- if performwork() throws an error we should releasethelock(lockkey) - finally.
Question: Why “releasethelock” if lockUntil is used?
Suppose a user grabs the lock for 1 hour for a 10-minute task and encounters an exception, nobody else can grab the lock for that hour.
Let's see some code
Add the below inside a Worker class with @Component and @EnableScheduling.
private static final Logger LOG = LoggerFactory.getLogger(Worker.class);
private final MongoTemplate mongoTemplate;
private final String lockedBy;
private final FindAndModifyOptions findAndModifyOptions;
@Autowired
public Worker(MongoTemplate mongoTemplate) {
this.mongoTemplate = mongoTemplate;
this.lockedBy = UUID.randomUUID().toString();
this.findAndModifyOptions = FindAndModifyOptions.options().returnNew(true).upsert(true);
}
Pseudocode Implementation
@Scheduled(fixedRate = 60_000)
void work() throws InterruptedException {
String lockkey = "anirudh";
boolean g = grabTheLock(lockkey);
if (g) {
// grabbed the lock, ready to perform some work.
try {
LOG.info("THE WORK!");
Thread.sleep(10_000);
} catch (Exception e) {
LOG.error("Exception occurred while executing task.", e);
} finally {
releaseTheLock(lockkey);
}
} else {
// could not grab the lock, cheers.
}
}
Grabbing the lock Implementation
protected boolean grabTheLock(String lockKey) {
boolean managedToGrabTheLock = false;
Update update = new Update();
Date lockedAt = new Date();
Date lockedUntil = DateUtils.addMinutes(lockedAt, 5);
Query query = Query.query(Criteria.where("lockKey").is(lockKey).and("lockedUntil").lt(lockedAt));
update.set("lockKey", lockKey);
update.set("lockedAt", lockedAt);
update.set("lockedUntil", lockedUntil);
update.set("lockedBy", lockedBy);
try {
MongoLock mongoLock = mongoTemplate.findAndModify(query, update, findAndModifyOptions, MongoLock.class);
LOG.info("Successfully grabbed the lock: {}", mongoLock);
managedToGrabTheLock = true;
} catch (DuplicateKeyException duplicateKeyException) {
// this exception is normal and is thrown when another instance has the lock and this instance is trying to insert with the lockKey.
LOG.info("Failed to grab the lock: {} at {} until {}.", lockKey, lockedAt, lockedUntil);
} catch (Exception e) {
LOG.error("Failed to grab the lock: {} at {} until {}.", lockKey, lockedAt, lockedUntil);
LOG.error(e.getMessage(), e);
}
return managedToGrabTheLock;
}
Releasing the lock implementation
private void releaseTheLock(String lockKey) {
Criteria criteria = Criteria.where("lockKey").is(lockKey)
.and("lockedBy").is(this.lockedBy);
Query query = Query.query(criteria);
MongoLock releasedLock = mongoTemplate.findAndRemove(query, MongoLock.class);
if (releasedLock != null) {
LOG.info("Successfully released the lock: {}", releasedLock);
} else {
LOG.error("Cannot release {} as it doesn't exist.", releasedLock);
}
}
Add annotate your SpringBootApplication like this.
@EnableScheduling
@SpringBootApplication
public class DistributedlocksfromscratchApplication {
public static void main(String[] args) {
SpringApplication.run(DistributedlocksfromscratchApplication.class, args);
}
}
MongoLock
@Document(value = "MongoLock")
public class MongoLock {
@Id
private String lockKey;
private Date lockedAt;
private Date lockedUntil;
private String lockedBy;
// Empty and a proper constructor.
// getters and setters
// to string
}
pom.xml - add the below.
<dependency>
<groupId>org.springframework.boot</groupId>
<artifactId>spring-boot-starter-data-mongodb</artifactId>
</dependency>
<dependency>
<groupId>org.apache.commons</groupId>
<artifactId>commons-lang3</artifactId>
</dependency>
application.properties
spring.data.mongodb.uri=mongodb://localhost:27017/test
docker-compose.yml file?
version: "3.9"
services:
mongo:
image: mongo
#restart: always
ports:
- 27017:27017
^ docker-compose up
Finally, we are done with the setup.
Let's run it :)
Iteration1:
One JVM, sample run:
2021-08-15 07:49:13.027 INFO 98862 --- [ scheduling-1] c.e.d.code.Worker : Successfully grabbed the lock: MongoLock[lockKey='anirudh', lockedAt=Sun Aug 15 07:49:13 EDT 2021, lockedUntil=Sun Aug 15 07:54:13 EDT 2021, lockedBy='5852d446-1922-47b8-98f3-1eb677524546']
2021-08-15 07:49:13.027 INFO 98862 --- [ scheduling-1] c.e.d.code.Worker : THE WORK!
2021-08-15 07:49:23.032 INFO 98862 --- [ scheduling-1] c.e.d.code.Worker : Successfully released the lock: MongoLock[lockKey='anirudh', lockedAt=Sun Aug 15 07:49:13 EDT 2021, lockedUntil=Sun Aug 15 07:54:13 EDT 2021, lockedBy='5852d446-1922-47b8-98f3-1eb677524546']
Line 1 says the lock will be held from 07:49:13 at most until 07:54:13 (5 mins).
To verify the above, add a breakpoint at “The work!” and run db.getCollection(‘MongoLock’).find({}). You will see an entry. Beautiful.
Iteration2:
To run multiple JVMs, please check the “Allow parallel run” on “Run/Debug Configurations” on Intellij.
Run it twice. Make sure the second one runs inside 10 seconds of the first (so that it fails).
JVM1
First run
2021-08-15 07:50:13.028 INFO 98862 --- [ scheduling-1] c.e.d.code.Worker : Successfully grabbed the lock: MongoLock[lockKey='anirudh', lockedAt=Sun Aug 15 07:50:13 EDT 2021, lockedUntil=Sun Aug 15 07:55:13 EDT 2021, lockedBy='5852d446-1922-47b8-98f3-1eb677524546']
2021-08-15 07:50:13.028 INFO 98862 --- [ scheduling-1] c.e.d.code.Worker : THE WORK!
2021-08-15 07:50:23.031 INFO 98862 --- [ scheduling-1] c.e.d.code.Worker : Successfully released the lock: MongoLock[lockKey='anirudh', lockedAt=Sun Aug 15 07:50:13 EDT 2021, lockedUntil=Sun Aug 15 07:55:13 EDT 2021, lockedBy='5852d446-1922-47b8-98f3-1eb677524546']
Second run
2021-08-15 07:51:13.025 INFO 98862 --- [ scheduling-1] c.e.d.code.Worker : Successfully grabbed the lock: MongoLock[lockKey='anirudh', lockedAt=Sun Aug 15 07:51:13 EDT 2021, lockedUntil=Sun Aug 15 07:56:13 EDT 2021, lockedBy='5852d446-1922-47b8-98f3-1eb677524546']
2021-08-15 07:51:13.026 INFO 98862 --- [ scheduling-1] c.e.d.code.Worker : THE WORK!
2021-08-15 07:51:23.034 INFO 98862 --- [ scheduling-1] c.e.d.code.Worker : Successfully released the lock: MongoLock[lockKey='anirudh', lockedAt=Sun Aug 15 07:51:13 EDT 2021, lockedUntil=Sun Aug 15 07:56:13 EDT 2021, lockedBy='5852d446-1922-47b8-98f3-1eb677524546']
JVM2
2021-08-15 07:50:20.271 INFO 98874 --- [ scheduling-1] c.e.d.code.Worker : Failed to grab the lock: anirudh at Sun Aug 15 07:50:20 EDT 2021 until Sun Aug 15 07:55:20 EDT 2021.
2021-08-15 07:51:20.272 INFO 98874 --- [ scheduling-1] c.e.d.code.Worker : Failed to grab the lock: anirudh at Sun Aug 15 07:51:20 EDT 2021 until Sun Aug 15 07:56:20 EDT 2021.
JVM1 acquired the lock at 07:50:13 and released it at 07:50:23.
JVM2 tried to acquire the lock at 07:50:20, but couldn't get it.
Next iteration (07:51:13) the same happened.
Hence, only one JVM acquires the lock at a time.
Conclusion
We learned how to implement distributed locking across multiple JVMs.
Gotchas
Please hold the lock for longer than expected.
Why?
Because while one instance is working on a shared state, another instance could grab the lock and work on the same shared state. Undesirable.
How to fix it?
Every instance should renew the mongolock every 10 seconds or so.
How to implement it?
Write another scheduled task that a) checks if I have a lock and proceeds to b) renews the lease.
Hope you enjoyed this newsletter!