diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockBase.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockBase.scala index 1d2f6a8e..4274729f 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockBase.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockBase.scala @@ -49,7 +49,7 @@ abstract class TokenLockBase(override val token: String) extends TokenLock { protected def tryAcquireGuardLock(retries: Int, thisTry: Int): Boolean - protected def releaseGuardLock(): Unit + protected def releaseGuardLock(evenNonOwned: Boolean): Unit protected def updateTicket(): Unit @@ -101,7 +101,7 @@ abstract class TokenLockBase(override val token: String) extends TokenLock { watcherThreadOpt.foreach(_.interrupt()) watcherThreadOpt = None try { - releaseGuardLock() + releaseGuardLock(evenNonOwned = false) } finally { JvmUtils.safeRemoveShutdownHook(shutdownHook) TokenLockRegistry.unregisterLock(this) @@ -136,7 +136,7 @@ abstract class TokenLockBase(override val token: String) extends TokenLock { if (wasAcquired) { watcherThreadOpt.foreach(_.interrupt()) watcherThreadOpt = None - releaseGuardLock() + releaseGuardLock(evenNonOwned = false) } } } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockDynamoDb.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockDynamoDb.scala index 10dafdc9..e3b7f0ba 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockDynamoDb.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockDynamoDb.scala @@ -78,7 +78,7 @@ class TokenLockDynamoDb( if (expires < now) { log.warn(s"Taking over expired ticket $escapedToken ($expires < $now)") - releaseGuardLock() + releaseGuardLock(evenNonOwned = true) tryAcquireGuardLock(retries - 1, thisTry + 1) } else { false @@ -105,25 +105,38 @@ class TokenLockDynamoDb( } /** Invoked from a synchronized block. */ - override def releaseGuardLock(): Unit = { + override def releaseGuardLock(evenNonOwned: Boolean): Unit = { try { val now = Instant.now() val nowEpoch = now.getEpochSecond val hardExpireTickets = now.minus(TICKETS_HARD_EXPIRE_DAYS, ChronoUnit.DAYS).getEpochSecond // Delete this ticket or any expired tickets - val deleteRequest = DeleteItemRequest.builder() + val deleteRequest = if (evenNonOwned) { + DeleteItemRequest.builder() .tableName(tableName) .key(Map( ATTR_TOKEN -> AttributeValue.builder().s(escapedToken).build() ).asJava) - .conditionExpression(s"$ATTR_OWNER = :jobOwner OR ($ATTR_EXPIRES < :now AND $ATTR_CREATED_AT < :hardExpire)") + .conditionExpression(s"$ATTR_EXPIRES < :now") .expressionAttributeValues(Map( - ":jobOwner" -> AttributeValue.builder().s(owner).build(), - ":now" -> AttributeValue.builder().n(nowEpoch.toString).build(), - ":hardExpire" -> AttributeValue.builder().n(hardExpireTickets.toString).build() + ":now" -> AttributeValue.builder().n(nowEpoch.toString).build() ).asJava) .build() + } else { + DeleteItemRequest.builder() + .tableName(tableName) + .key(Map( + ATTR_TOKEN -> AttributeValue.builder().s(escapedToken).build() + ).asJava) + .conditionExpression(s"$ATTR_OWNER = :jobOwner OR ($ATTR_EXPIRES < :now AND $ATTR_CREATED_AT < :hardExpire)") + .expressionAttributeValues(Map( + ":jobOwner" -> AttributeValue.builder().s(owner).build(), + ":now" -> AttributeValue.builder().n(nowEpoch.toString).build(), + ":hardExpire" -> AttributeValue.builder().n(hardExpireTickets.toString).build() + ).asJava) + .build() + } try { dynamoDbClient.deleteItem(deleteRequest) diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockHadoopPath.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockHadoopPath.scala index a0d032b0..5db82866 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockHadoopPath.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockHadoopPath.scala @@ -40,7 +40,7 @@ class TokenLockHadoopPath(token: String, } /** Invoked from a synchronized block. */ - override def releaseGuardLock(): Unit = { + override def releaseGuardLock(evenNonOwned: Boolean): Unit = { fileGuardOpt.foreach { fileGuard => fsUtils.deleteFile(fileGuard) fileGuardOpt = None diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockJdbc.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockJdbc.scala index 2f30ff4e..0579f3fe 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockJdbc.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockJdbc.scala @@ -56,9 +56,10 @@ class TokenLockJdbc(token: String, db: Database, slickProfile: JdbcProfile) exte val now = Instant.now().getEpochSecond if (expires < now) { log.warn(s"Taking over expired ticket $escapedToken ($expires < $now)") - releaseGuardLock() + releaseExpiredGuardLock(now) tryAcquireGuardLock(retries - 1, thisTry + 1) } else { + log.warn(s"The ticket for $escapedToken is still valid ($expires >= $now)") false } } @@ -93,19 +94,40 @@ class TokenLockJdbc(token: String, db: Database, slickProfile: JdbcProfile) exte } /** Invoked from a synchronized block. */ - override def releaseGuardLock(): Unit = { + override def releaseGuardLock(evenNonOwned: Boolean): Unit = { try { val now = Instant.now() val nowEpoch = now.getEpochSecond val hardExpireTickets = now.minus(TICKETS_HARD_EXPIRE_DAYS, ChronoUnit.DAYS).getEpochSecond + + if (evenNonOwned) { + slickUtils.executeAction(db, lockTicketTable.records.filter(ticket => ticket.token === escapedToken).delete) + } else { + slickUtils.executeAction( + db, + lockTicketTable.records + .filter(ticket => (ticket.token === escapedToken && ticket.owner === owner) || + (ticket.createdAt.isDefined && ticket.createdAt < hardExpireTickets && ticket.expires < nowEpoch)).delete + ) + } + } catch { + case NonFatal(ex) => log.error(s"An error occurred when trying to release the lock: $escapedToken.", ex) + } + } + + /** + * Invoked from a synchronized block. + * Removes the ticket only when both the token matches and the ticket is still expired. + */ + private def releaseExpiredGuardLock(now: Long): Unit = { + try { slickUtils.executeAction( db, lockTicketTable.records - .filter(ticket => (ticket.token === escapedToken && ticket.owner === owner) || - (ticket.createdAt.isDefined && ticket.createdAt < hardExpireTickets && ticket.expires < nowEpoch)).delete + .filter(ticket => ticket.token === escapedToken && ticket.expires < now).delete ) } catch { - case NonFatal(ex) => log.error(s"An error occurred when trying to release the lock: $escapedToken.", ex) + case NonFatal(ex) => log.error(s"An error occurred when trying to release the expired lock: $escapedToken.", ex) } } diff --git a/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockMongoDb.scala b/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockMongoDb.scala index fcb99286..fea6b7ca 100644 --- a/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockMongoDb.scala +++ b/pramen/core/src/main/scala/za/co/absa/pramen/core/lock/TokenLockMongoDb.scala @@ -67,7 +67,7 @@ class TokenLockMongoDb(token: String, val now = Instant.now().getEpochSecond if (expires < now) { log.warn(s"Taking over expired ticket $escapedToken ($expires < $now)") - releaseGuardLock() + releaseGuardLock(evenNonOwned = true) tryAcquireGuardLock(retries - 1, thisTry + 1) true } else { @@ -95,7 +95,7 @@ class TokenLockMongoDb(token: String, } /** Invoked from a synchronized block. */ - override def releaseGuardLock(): Unit = { + override def releaseGuardLock(evenNonOwned: Boolean): Unit = { try { val c = getCollection log.debug(s"Delete token $escapedToken") diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/lock/TokenLockJdbcSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/lock/TokenLockJdbcSuite.scala index 507f1d6c..e910b209 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/lock/TokenLockJdbcSuite.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/lock/TokenLockJdbcSuite.scala @@ -24,7 +24,7 @@ import za.co.absa.pramen.core.fixtures.RelationalDbFixture import za.co.absa.pramen.core.lock.{TokenLockBase, TokenLockJdbc, TokenLockRegistry} import za.co.absa.pramen.core.rdb.{PramenDb, RdbJdbc} import za.co.absa.pramen.core.reader.model.JdbcConfig -import za.co.absa.pramen.core.utils.UsingUtils +import za.co.absa.pramen.core.utils.{SlickUtils, UsingUtils} import scala.concurrent.duration._ @@ -82,6 +82,23 @@ class TokenLockJdbcSuite extends AnyWordSpec with RelationalDbFixture with Befor lock2.release() } + "allow releasing locks for other owners if requested" in { + val lock1 = getLock("token1") + val lock2 = getLock("token2") + + assert(lock1.tryAcquire()) + assert(lock2.tryAcquire()) + + lock1.asInstanceOf[TokenLockJdbc].releaseGuardLock(evenNonOwned = false) + lock2.asInstanceOf[TokenLockJdbc].releaseGuardLock(evenNonOwned = true) + + val slickUtils = new SlickUtils(pramenDb.slickProfile) + + val recordCount = slickUtils.executeCount(pramenDb.slickDb, pramenDb.lockTicketTable.records.length) + + assert(recordCount == 0) + } + "lock pramen should constantly update lock ticket" in { val lock1 = new TokenLockJdbc("token1", pramenDb.slickDb, pramenDb.slickProfile) { override val tokenExpiresSeconds = 2L diff --git a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/lock/TokenLockSQLiteSuite.scala b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/lock/TokenLockSQLiteSuite.scala index 679888da..6c175645 100644 --- a/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/lock/TokenLockSQLiteSuite.scala +++ b/pramen/core/src/test/scala/za/co/absa/pramen/core/tests/lock/TokenLockSQLiteSuite.scala @@ -24,6 +24,7 @@ import za.co.absa.pramen.core.fixtures.TempDirFixture import za.co.absa.pramen.core.lock.TokenLockJdbc import za.co.absa.pramen.core.rdb.PramenDb import za.co.absa.pramen.core.reader.model.JdbcConfig +import za.co.absa.pramen.core.utils.SlickUtils import java.io.File @@ -93,6 +94,23 @@ class TokenLockSQLiteSuite extends AnyWordSpec with BeforeAndAfter with BeforeA lock2.release() } + "allow releasing locks for other owners if requested" in { + val lock1 = getLock("token1") + val lock2 = getLock("token2") + + assert(lock1.tryAcquire()) + assert(lock2.tryAcquire()) + + lock1.asInstanceOf[TokenLockJdbc].releaseGuardLock(evenNonOwned = false) + lock2.asInstanceOf[TokenLockJdbc].releaseGuardLock(evenNonOwned = true) + + val slickUtils = new SlickUtils(pramenDb.slickProfile) + + val recordCount = slickUtils.executeCount(pramenDb.slickDb, pramenDb.lockTicketTable.records.length) + + assert(recordCount == 0) + } + "lock pramen should constantly update lock ticket" in { val lock1 = new TokenLockJdbc("token1", pramenDb.slickDb, pramenDb.slickProfile) { override val tokenExpiresSeconds = 2L