anishshri-db commented on code in PR #50595:
URL: https://github.com/apache/spark/pull/50595#discussion_r2064407173
##########
sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/StateStore.scala:
##########
@@ -1173,6 +1203,114 @@ object StateStore extends Logging {
}
}
+ private def processQueuedProviders(storeConf: StateStoreConf): Unit = {
+ // Get up to 5 queued providers at a time to avoid blocking maintenance
for too long
+ val maxProvidersToProcess = 5
+ val providersToProcess = new
mutable.ArrayBuffer[StateStoreProviderId](maxProvidersToProcess)
+
+ for (_ <- 1 to maxProvidersToProcess if !partitionsForMaintenance.isEmpty)
{
+ val id = partitionsForMaintenance.poll()
+ if (id != null) {
+ providersToProcess += id
+ }
+ }
+
+ providersToProcess.foreach { id =>
+ loadedProviders.synchronized {
+ loadedProviders.get(id)
+ }.foreach { provider =>
+ if (processThisPartition(id)) {
+ logInfo(s"Processing previously queued provider $id for maintenance")
+ submitMaintenanceWorkForProvider(id, provider, storeConf)
+ } else {
+ // If we couldn't get the lock, put it back in the queue
+ partitionsForMaintenance.add(id)
+ }
+ }
+ }
+ }
+
+ /**
+ * Submits maintenance work for a provider to the maintenance thread pool.
+ *
+ * @param id The StateStore provider ID to perform maintenance on
+ * @param provider The StateStore provider instance
+ * @param submittedFromTaskThread If true, provider was already removed from
+ * loadedProviders. If false, we must already
+ * have acquired the lock to process this
partition.
+ */
+ private def submitMaintenanceWorkForProvider(
+ id: StateStoreProviderId,
+ provider: StateStoreProvider,
+ storeConf: StateStoreConf,
+ submittedFromTaskThread: Boolean = false): Unit = {
+ maintenanceThreadPool.execute(() => {
+ val startTime = System.currentTimeMillis()
+ if (submittedFromTaskThread) {
+ // If provider is already removed from loadedProviders (which can
happen when a task thread
+ // triggers unloading of an old provider), we MUST process this
partition to
+ // close it properly.
+ // We block here until we can acquire the lock for this partition,
waiting for any
+ // possible ongoing maintenance on this partition to complete first.
+ val ableToProcess = awaitProcessThisPartition(id, storeConf)
+ if (!ableToProcess) {
+ // The provider has been added to the providersToClose queue in
awaitProcessThisPartition
+ // so we can return early
+ return
+ }
+ }
+ val awaitingPartitionDuration = System.currentTimeMillis() - startTime
+ try {
+ provider.doMaintenance()
+ // If submittedFromTaskThread is false, we don't need to verify
+ // with the coordinator as we know it definitely should be unloaded.
+ if (submittedFromTaskThread) {
+ unload(id, Some(provider))
Review Comment:
nit: indent is off ?
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org
---------------------------------------------------------------------
To unsubscribe, e-mail: reviews-unsubscr...@spark.apache.org
For additional commands, e-mail: reviews-h...@spark.apache.org
