prevent duplicate queueing in the prio semaphore

sql-plugin/src/main/scala/com/nvidia/spark/rapids/GpuSemaphore.scala

@@ -288,7 +288,7 @@ private final class SemaphoreTaskInfo() extends Logging {
     } else {
       if (blockedThreads.size() == 0) {
         // No other threads for this task are waiting, so we might be able to grab this directly
-        val ret = semaphore.tryAcquire(numPermits)
+        val ret = semaphore.tryAcquire(numPermits, lastHeld)
         if (ret) {
           hasSemaphore = true
           activeThreads.add(t)

sql-plugin/src/main/scala/com/nvidia/spark/rapids/PrioritySemaphore.scala

@@ -32,22 +32,27 @@ class PrioritySemaphore[T](val maxPermits: Int)(implicit ordering: Ordering[T])
   private val lock = new ReentrantLock()
   private var occupiedSlots: Int = 0
 
-  private case class ThreadInfo(priority: T, condition: Condition)
+  private case class ThreadInfo(priority: T, condition: Condition, numPermits: Int) {
+    var signaled: Boolean = false
+    var interrupted: Boolean = false
+  }
 
   // We expect a relatively small number of threads to be contending for this lock at any given
   // time, therefore we are not concerned with the insertion/removal time complexity.
   private val waitingQueue: PriorityQueue[ThreadInfo] = PriorityQueue()(Ordering.by(_.priority))
 
   def this()(implicit ordering: Ordering[T]) = this(PrioritySemaphore.DEFAULT_MAX_PERMITS)(ordering)
 
-  def tryAcquire(numPermits: Int): Boolean = {
+  def tryAcquire(numPermits: Int, priority: T): Boolean = {
     lock.lock()
     try {
-      if (canAcquire(numPermits)) {
+      if (waitingQueue.nonEmpty && ordering.gt(waitingQueue.head.priority, priority)) {
+        false
+      } else if (!canAcquire(numPermits)) {
+        false
+      } else {
         commitAcquire(numPermits)
         true
-      } else {
-        false
       }
     } finally {
       lock.unlock()
@@ -56,17 +61,24 @@ class PrioritySemaphore[T](val maxPermits: Int)(implicit ordering: Ordering[T])
 
   def acquire(numPermits: Int, priority: T): Unit = {
     lock.lock()
+    val condition = lock.newCondition()
+    val info = ThreadInfo(priority, condition, numPermits)
     try {
-      val condition = lock.newCondition()
-      while (!canAcquire(numPermits)) {
-        waitingQueue.enqueue(ThreadInfo(priority, condition))
-        condition.await()
+      if (canAcquire(numPermits)) {
+        commitAcquire(numPermits)
+      } else {
+        waitingQueue.enqueue(info)
+        while (!info.signaled) {
+          info.condition.await()
+        }
       }
-      commitAcquire(numPermits)
-
+    } catch {
+      case _: InterruptedException =>
+        info.interrupted = true
     } finally {
       lock.unlock()
-    }}
+    }
+  }
 
   private def commitAcquire(numPermits: Int): Unit = {
     occupiedSlots += numPermits
@@ -76,9 +88,21 @@ class PrioritySemaphore[T](val maxPermits: Int)(implicit ordering: Ordering[T])
     lock.lock()
     try {
       occupiedSlots -= numPermits
-      if (waitingQueue.nonEmpty) {
-        val nextThread = waitingQueue.dequeue()
-        nextThread.condition.signal()
+      // acquire and wakeup for all threads that now have enough permits
+      var done = false
+      while (!done && waitingQueue.nonEmpty) {
+        val nextThread = waitingQueue.head
+        if (nextThread.interrupted) {
+          waitingQueue.dequeue()
+        } else if (canAcquire(nextThread.numPermits)) {
+          val popped = waitingQueue.dequeue()
+          assert(popped eq nextThread)
+          commitAcquire(nextThread.numPermits)
+          nextThread.signaled = true
+          nextThread.condition.signal()
+        } else {
+          done = true
+        }
       }
     } finally {
       lock.unlock()

tests/src/test/scala/com/nvidia/spark/rapids/PrioritySemaphoreSuite.scala

@@ -28,16 +28,16 @@ class PrioritySemaphoreSuite extends AnyFunSuite {
   test("tryAcquire should return true if permits are available") {
     val semaphore = new TestPrioritySemaphore(10)
 
-    assert(semaphore.tryAcquire(5))
-    assert(semaphore.tryAcquire(3))
-    assert(semaphore.tryAcquire(2))
-    assert(!semaphore.tryAcquire(1))
+    assert(semaphore.tryAcquire(5, 0))
+    assert(semaphore.tryAcquire(3, 0))
+    assert(semaphore.tryAcquire(2, 0))
+    assert(!semaphore.tryAcquire(1, 0))
   }
 
   test("acquire and release should work correctly") {
     val semaphore = new TestPrioritySemaphore(1)
 
-    assert(semaphore.tryAcquire(1))
+    assert(semaphore.tryAcquire(1, 0))
 
     val latch = new CountDownLatch(1)
     val t = new Thread(() => {