Skip to content
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

Closed
wants to merge 4 commits into from
Closed
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -124,12 +125,46 @@ 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
/**
* Runs the `task` periodically and automatically cancels it if there is an exception. `onError`
* will be called when an exception happens.
*/
class MaintenanceTask(periodMs: Long, task: => Unit, onError: => Unit) {
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you add docs.

Copy link
Contributor

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.

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(
Expand Down Expand Up @@ -162,15 +197,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")
Expand All @@ -179,14 +214,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 && !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")
}
}
Expand All @@ -198,21 +233,20 @@ object StateStore extends Logging {
private def doMaintenance(): Unit = {
logDebug("Doing maintenance")
if (SparkEnv.get == null) {
Copy link
Member Author

Choose a reason for hiding this comment

The 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 SparkEnv.get == null, the following stop may cancel a new valid task. However, I think that won't happen in practice, so don't fix it.

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
}
}
}
Expand All @@ -238,7 +272,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) {
Expand Down