adding lock

Author: ericm-db

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala

@@ -856,6 +856,9 @@ object StateStore extends Logging {
 
   private val maintenanceThreadPoolLock = new Object
 
+  private val providersQueueLock = new Object
+
+  @GuardedBy("providersQueueLock")
   private val unloadedProvidersToClose =
     new ConcurrentLinkedQueue[(StateStoreProviderId, StateStoreProvider)]
 
@@ -1162,17 +1165,19 @@ object StateStore extends Logging {
   // Wait until this partition can be processed
   private def awaitProcessThisPartition(
       id: StateStoreProviderId,
-      timeoutMs: Long
-  ): Boolean = maintenanceThreadPoolLock synchronized  {
-    val endTime = System.currentTimeMillis() + timeoutMs
+      timeoutMs: Long): Boolean = maintenanceThreadPoolLock synchronized  {
+    val startTime = System.currentTimeMillis()
+    val endTime = startTime + timeoutMs
 
     // If immediate processing fails, wait with timeout
     var canProcessThisPartition = processThisPartition(id)
     while (!canProcessThisPartition && System.currentTimeMillis() < endTime) {
       canProcessThisPartition = processThisPartition(id)
       maintenanceThreadPoolLock.wait(timeoutMs)
     }
-
+    val elapsedTime = System.currentTimeMillis() - startTime
+    logInfo(log"Waited for ${MDC(LogKeys.TOTAL_TIME, elapsedTime)} ms to be able to process " +
+      log"maintenance for partition ${MDC(LogKeys.STATE_STORE_PROVIDER_ID, id)}")
     canProcessThisPartition
   }
 
@@ -1189,20 +1194,34 @@ object StateStore extends Logging {
     }
 
     // Providers that couldn't be processed now and need to be added back to the queue
-    val providersToRequeue = new ArrayBuffer[(StateStoreProviderId, StateStoreProvider)]()
-
-    while (!unloadedProvidersToClose.isEmpty) {
-      val (providerId, provider) = unloadedProvidersToClose.poll()
+    val providersToRequeue = ArrayBuffer.empty[(StateStoreProviderId, StateStoreProvider)]
+
+    // Create a temporary list and drain the concurrent queue into it under a lock
+    val tempList = providersQueueLock synchronized {
+      val items = ArrayBuffer.empty[(StateStoreProviderId, StateStoreProvider)]
+      while (!unloadedProvidersToClose.isEmpty) {
+        val item = unloadedProvidersToClose.poll()
+        items += item
+      }
+      items
+    }
 
+    // Process all items in the temporary list
+    tempList.foreach { case (providerId, provider) =>
       if (processThisPartition(providerId)) {
         submitMaintenanceWorkForProvider(
           providerId, provider, storeConf, MaintenanceTaskType.FromMaintenanceQueue)
       } else {
-        providersToRequeue += (providerId, provider)
+        providersToRequeue += ((providerId, provider))
       }
     }
 
-    providersToRequeue.foreach(unloadedProvidersToClose.offer)
+    providersQueueLock synchronized {
+      // Add providers that couldn't be processed back to the original queue
+      providersToRequeue.foreach { providerPair =>
+        unloadedProvidersToClose.offer(providerPair)
+      }
+    }
 
     loadedProviders.synchronized {
       loadedProviders.toSeq
@@ -1239,7 +1258,9 @@ object StateStore extends Logging {
           val ableToProcessNow = awaitProcessThisPartition(id, timeoutMs)
           if (!ableToProcessNow) {
             // Add to queue for later processing if we can't process now
-            unloadedProvidersToClose.add((id, provider))
+            providersQueueLock synchronized {
+              unloadedProvidersToClose.offer((id, provider))
+            }
           }
           ableToProcessNow
 
@@ -1249,6 +1270,8 @@ object StateStore extends Logging {
           true
 
         case FromLoadedProviders =>
+          // Provider from loadedProviders can be processed immediately
+          // as it's in maintenancePartitions
           true
       }