-
Notifications
You must be signed in to change notification settings - Fork 28.4k
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
[SPARK-19267][SS]Fix a race condition when stopping StateStore #16627
Changes from 3 commits
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -18,6 +18,7 @@ | |
package org.apache.spark.sql.execution.streaming.state | ||
|
||
import java.util.concurrent.{ScheduledFuture, TimeUnit} | ||
import javax.annotation.concurrent.GuardedBy | ||
|
||
import scala.collection.mutable | ||
import scala.util.control.NonFatal | ||
|
@@ -124,12 +125,42 @@ object StateStore extends Logging { | |
val MAINTENANCE_INTERVAL_CONFIG = "spark.sql.streaming.stateStore.maintenanceInterval" | ||
val MAINTENANCE_INTERVAL_DEFAULT_SECS = 60 | ||
|
||
@GuardedBy("loadedProviders") | ||
private val loadedProviders = new mutable.HashMap[StateStoreId, StateStoreProvider]() | ||
private val maintenanceTaskExecutor = | ||
ThreadUtils.newDaemonSingleThreadScheduledExecutor("state-store-maintenance-task") | ||
|
||
@volatile private var maintenanceTask: ScheduledFuture[_] = null | ||
@volatile private var _coordRef: StateStoreCoordinatorRef = null | ||
class MaintenanceTask(periodMs: Long, task: => Unit, onError: => Unit) { | ||
private val executor = | ||
ThreadUtils.newDaemonSingleThreadScheduledExecutor("state-store-maintenance-task") | ||
|
||
private val runnable = new Runnable { | ||
override def run(): Unit = { | ||
try { | ||
task | ||
} catch { | ||
case NonFatal(e) => | ||
logWarning("Error running maintenance thread", e) | ||
onError | ||
throw e | ||
} | ||
} | ||
} | ||
|
||
private val future: ScheduledFuture[_] = executor.scheduleAtFixedRate( | ||
runnable, periodMs, periodMs, TimeUnit.MILLISECONDS) | ||
|
||
def stop(): Unit = { | ||
future.cancel(false) | ||
executor.shutdown() | ||
} | ||
|
||
def isRunning: Boolean = !future.isDone | ||
} | ||
|
||
@GuardedBy("loadedProviders") | ||
private var maintenanceTask: MaintenanceTask = null | ||
|
||
@GuardedBy("loadedProviders") | ||
private var _coordRef: StateStoreCoordinatorRef = null | ||
|
||
/** Get or create a store associated with the id. */ | ||
def get( | ||
|
@@ -162,15 +193,15 @@ object StateStore extends Logging { | |
} | ||
|
||
def isMaintenanceRunning: Boolean = loadedProviders.synchronized { | ||
maintenanceTask != null | ||
maintenanceTask != null && maintenanceTask.isRunning | ||
} | ||
|
||
/** Unload and stop all state store providers */ | ||
def stop(): Unit = loadedProviders.synchronized { | ||
loadedProviders.clear() | ||
_coordRef = null | ||
if (maintenanceTask != null) { | ||
maintenanceTask.cancel(false) | ||
maintenanceTask.stop() | ||
maintenanceTask = null | ||
} | ||
logInfo("StateStore stopped") | ||
|
@@ -179,14 +210,14 @@ object StateStore extends Logging { | |
/** Start the periodic maintenance task if not already started and if Spark active */ | ||
private def startMaintenanceIfNeeded(): Unit = loadedProviders.synchronized { | ||
val env = SparkEnv.get | ||
if (maintenanceTask == null && env != null) { | ||
if (env != null && (maintenanceTask == null || !maintenanceTask.isRunning)) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Can you replace this with the method isMaintenanceRunning? |
||
val periodMs = env.conf.getTimeAsMs( | ||
MAINTENANCE_INTERVAL_CONFIG, s"${MAINTENANCE_INTERVAL_DEFAULT_SECS}s") | ||
val runnable = new Runnable { | ||
override def run(): Unit = { doMaintenance() } | ||
} | ||
maintenanceTask = maintenanceTaskExecutor.scheduleAtFixedRate( | ||
runnable, periodMs, periodMs, TimeUnit.MILLISECONDS) | ||
maintenanceTask = new MaintenanceTask( | ||
periodMs, | ||
task = { doMaintenance() }, | ||
onError = { loadedProviders.synchronized { loadedProviders.clear() } } | ||
) | ||
logInfo("State Store maintenance task started") | ||
} | ||
} | ||
|
@@ -198,21 +229,20 @@ object StateStore extends Logging { | |
private def doMaintenance(): Unit = { | ||
logDebug("Doing maintenance") | ||
if (SparkEnv.get == null) { | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. There is another potential issue here: if a SparkContext is created after checking |
||
stop() | ||
} else { | ||
loadedProviders.synchronized { loadedProviders.toSeq }.foreach { case (id, provider) => | ||
try { | ||
if (verifyIfStoreInstanceActive(id)) { | ||
provider.doMaintenance() | ||
} else { | ||
unload(id) | ||
logInfo(s"Unloaded $provider") | ||
} | ||
} catch { | ||
case NonFatal(e) => | ||
logWarning(s"Error managing $provider, stopping management thread") | ||
stop() | ||
throw new IllegalStateException("SparkEnv not active, cannot do maintenance on StateStores") | ||
} | ||
loadedProviders.synchronized { loadedProviders.toSeq }.foreach { case (id, provider) => | ||
try { | ||
if (verifyIfStoreInstanceActive(id)) { | ||
provider.doMaintenance() | ||
} else { | ||
unload(id) | ||
logInfo(s"Unloaded $provider") | ||
} | ||
} catch { | ||
case NonFatal(e) => | ||
logWarning(s"Error managing $provider, stopping management thread") | ||
throw e | ||
} | ||
} | ||
} | ||
|
@@ -238,7 +268,7 @@ object StateStore extends Logging { | |
} | ||
} | ||
|
||
private def coordinatorRef: Option[StateStoreCoordinatorRef] = synchronized { | ||
private def coordinatorRef: Option[StateStoreCoordinatorRef] = loadedProviders.synchronized { | ||
val env = SparkEnv.get | ||
if (env != null) { | ||
if (_coordRef == null) { | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Can you add docs.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
you should mention the properties of this class. that it automatically cancels the periodic task if there is an exception. and what is onError for.