Skip to content

[SPARK-3137][Core]Replace the global TorrentBroadcast lock with fine grained KeyLock #25612

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.

Sign up for GitHub

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

Closed
zsxwing wants to merge 5 commits into from
Closed

[SPARK-3137][Core]Replace the global TorrentBroadcast lock with fine grained KeyLock #25612

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
290a3e5
Replace the global TorrentBroadcast lock with fine grained locks
zsxwing Aug 23, 2019
2cf1826
add KeyLock so we can reuse this pattern later
zsxwing Aug 29, 2019
844be5c
update comment
zsxwing Aug 29, 2019
2e9639c
test
zsxwing Aug 30, 2019
9b59ca2
address nits
zsxwing Sep 3, 2019
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
6 changes: 5 additions & 1 deletion core/src/main/scala/org/apache/spark/broadcast/BroadcastManager.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,7 @@

package org.apache.spark.broadcast

import java.util.Collections
import java.util.concurrent.atomic.AtomicLong

import scala.reflect.ClassTag
Expand Down Expand Up @@ -56,7 +57,10 @@ private[spark] class BroadcastManager(
private val nextBroadcastId = new AtomicLong(0)

private[broadcast] val cachedValues = {

@srowen srowen Aug 30, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

You don't need the braces here, but I wouldn't bother unless you're changing this file again

new ReferenceMap(AbstractReferenceMap.HARD, AbstractReferenceMap.WEAK)
Collections.synchronizedMap(
new ReferenceMap(AbstractReferenceMap.HARD, AbstractReferenceMap.WEAK)
.asInstanceOf[java.util.Map[Any, Any]]

@zsxwing zsxwing Aug 28, 2019

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

This one now needs protection as the global lock is removed. Casting it because Scala compiler cannot figure it out.

@juliuszsompolski juliuszsompolski Sep 3, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

nit: would it make compiler happy and be nicer if the val was typed (val cachedValues: JMap[Any, Any]) instead of doing asInstanceOf?

@zsxwing zsxwing Sep 3, 2019

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

No. The compiler doesn't like Java classes :P

Comment thread
srowen marked this conversation as resolved.
)
}

def newBroadcast[T: ClassTag](value_ : T, isLocal: Boolean): Broadcast[T] = {
Expand Down
49 changes: 44 additions & 5 deletions core/src/main/scala/org/apache/spark/broadcast/TorrentBroadcast.scala
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,7 @@ package org.apache.spark.broadcast
import java.io._
import java.lang.ref.SoftReference
import java.nio.ByteBuffer
import java.util.concurrent.ConcurrentHashMap
import java.util.zip.Adler32

import scala.collection.JavaConverters._
Expand Down Expand Up @@ -167,7 +168,7 @@ private[spark] class TorrentBroadcast[T: ClassTag](obj: T, id: Long)
bm.getLocalBytes(pieceId) match {
case Some(block) =>
blocks(pid) = block
releaseLock(pieceId)
releaseBlockManagerLock(pieceId)
case None =>
bm.getRemoteBytes(pieceId) match {
case Some(b) =>
Expand Down Expand Up @@ -215,8 +216,10 @@ private[spark] class TorrentBroadcast[T: ClassTag](obj: T, id: Long)
}

private def readBroadcastBlock(): T = Utils.tryOrIOException {
val broadcastCache = SparkEnv.get.broadcastManager.cachedValues
broadcastCache.synchronized {
TorrentBroadcast.withTorrentBroadcastLock(broadcastId) {
// As we just lock based on `broadcastId`, whenever using `broadcastCache`, we should just

@juliuszsompolski juliuszsompolski Sep 3, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

nit: just -> only

// touch `broadcastId`.
val broadcastCache = SparkEnv.get.broadcastManager.cachedValues

Option(broadcastCache.get(broadcastId)).map(_.asInstanceOf[T]).getOrElse {
setConf(SparkEnv.get.conf)
Expand All @@ -225,7 +228,7 @@ private[spark] class TorrentBroadcast[T: ClassTag](obj: T, id: Long)
case Some(blockResult) =>
if (blockResult.data.hasNext) {
val x = blockResult.data.next().asInstanceOf[T]
releaseLock(broadcastId)
releaseBlockManagerLock(broadcastId)

if (x != null) {
broadcastCache.put(broadcastId, x)
Expand Down Expand Up @@ -270,7 +273,7 @@ private[spark] class TorrentBroadcast[T: ClassTag](obj: T, id: Long)
* If running in a task, register the given block's locks for release upon task completion.
* Otherwise, if not running in a task then immediately release the lock.
*/
private def releaseLock(blockId: BlockId): Unit = {
private def releaseBlockManagerLock(blockId: BlockId): Unit = {

@zsxwing zsxwing Aug 28, 2019

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Rename this method to avoid confusing.

val blockManager = SparkEnv.get.blockManager
Option(TaskContext.get()) match {
case Some(taskContext) =>
Expand All @@ -290,6 +293,42 @@ private[spark] class TorrentBroadcast[T: ClassTag](obj: T, id: Long)

private object TorrentBroadcast extends Logging {

/** Locks to ensure there is only one thread fetching the same [[TorrentBroadcast]] block. */
private val torrentBroadcastLock = new ConcurrentHashMap[BroadcastBlockId, AnyRef]()

/** Acquire a lock for fetching a [[TorrentBroadcast]] block. */
private def acquireTorrentBroadcastLock(broadcastId: BroadcastBlockId): Unit = {
while (true) {
val lock = torrentBroadcastLock.putIfAbsent(broadcastId, new Object)
if (lock == null) return

@srowen srowen Aug 28, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Can this happen?

@zsxwing zsxwing Aug 28, 2019

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Yep when no such key in ConcurrentHashMap.

@srowen srowen Aug 28, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Right OK, but then it isn't locking on anything the first time, how does this work?

@zsxwing zsxwing Aug 28, 2019 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

  • The caller adding an object to ConcurrentHashMap successfully wins the lock and acquireTorrentBroadcastLock will return. Others should wait until they add an object to ConcurrentHashMap successfully.
  • In releaseTorrentBroadcastLock, it will remove the object from ConcurrentHashMap and wake up threads waiting for this object to retry.

@srowen srowen Aug 28, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Sorry I get it now. The Object isn't really the lock. It's the presence of absence of the Object. The Object is used to queue up waiting callers.

OK but this seems simpler if you computeIfAbsent a ReentrantLock and just use the lock as a lock directly. It manages all the details of allowing one caller, waking up waiters, etc. I'd imagine it's more straightforward to understand?

@zsxwing zsxwing Aug 28, 2019

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

OK but this seems simpler if you computeIfAbsent a ReentrantLock and just use the lock as a lock directly. It manages all the details of allowing one caller, waking up waiters, etc. I'd imagine it's more straightforward to understand?

The unlock code path is pretty tricky. Needs to unlock and remove from the map (the removal will allow a thread to put a new lock). I can try to make these two actions atomic by adding a global lock. But that's more complicated than the current way.

@srowen srowen Aug 28, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

Isn't this all there is to it, more or less?

val lock = torrentBroadcastLock.computeIfAbsent(broadcastId, new ReentrantLock())
lock.lock()
try {
  func
} finally {
  lock.unlock()
}

The downside I suppose is that you have a map of locks that keeps growing, but how many broadcast IDs would there be? vs a relatively small lock object for each.

@zsxwing zsxwing Aug 29, 2019 •
edited
Loading

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

The number of broadcast IDs is significant enough for a long running application. Even if the user doesn't use broadcast, we will still use it internally, for example, task binary is a broadcast object, a Spark SQL query may use broadcast join...

@srowen srowen Aug 29, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

OK, makes sense. Well, could still make sense to use a Condition but it wouldn't change much.

lock.synchronized {
while (torrentBroadcastLock.get(broadcastId) eq lock) {
Comment thread
srowen marked this conversation as resolved.
Outdated
lock.wait()

@srowen srowen Aug 28, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

How about just using a Java Lock implementation rather than re-roll the signalling?

@zsxwing zsxwing Aug 28, 2019

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

Using a <ReentrantLock, Condition> pair? The codes would be more complicated.

@srowen srowen Aug 28, 2019

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

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

No, just a ReentrantLock. I think you don't need to deal with spurious wakeups then either

@zsxwing zsxwing Aug 28, 2019

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

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

I just tried ReentrantLock but didn't find a simple way to implement the same function.

@juliuszsompolski juliuszsompolski Aug 29, 2019

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

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

Why not just stripe over a fixed number of Reentrant locks? This should reduce contention enough (and potential collissions can be controlled with numLocks) is easy to understand and doesn't reimplement its own lock queues. It's also fully interruptible, while in the current implementation the threads that are queueing for the lock are not interruptible.

private val numLocks = 100
private val locks = (1 to numLocks).map(_ => new java.util.concurrent.locks.ReentrantLock)

private def withTorrentBroadcastLock[T](broadcastId: BroadcastBlockId)(func: => T): T = {
  val lock = locks(broadcastId.broadcastId % numLocks)
  lock.lockInterruptibly()
  try {
    func
  } finally {
    lock.unlock()
  }
}

}
}
}
}

/** Release the lock for a [[TorrentBroadcast]] block. */
private def releaseTorrentBroadcastLock(broadcastId: BroadcastBlockId): Unit = {
val lock = torrentBroadcastLock.remove(broadcastId)
if (lock != null) {
lock.synchronized {
lock.notifyAll()
}
}
}

/** Acquire and hold the lock to access `broadcastId` when calling `func`. */
private def withTorrentBroadcastLock[T](broadcastId: BroadcastBlockId)(func: => T): T = {
acquireTorrentBroadcastLock(broadcastId)
try {
func
} finally {
releaseTorrentBroadcastLock(broadcastId)
}
}

def blockifyObject[T: ClassTag](
obj: T,
blockSize: Int,
Expand Down