Spark的TorrentBroadcast：實現

devos發表於2015-08-19
SparkAST

依據Spark 1.4版

序列化和反序列化

前邊提到，TorrentBroadcast的關鍵就在於特殊的序列化和反序列化設定。1.1版的TorrentBroadcast實現了自己的readObject和writeObject方法，但是1.4.1版的TorrentBroadcast沒有實現自己的readObject方法，那麼它是如何進行序列化和反序列化的呢？

// obj就是被廣播的物件
 private val numBlocks: Int = writeBlocks(obj)

  override protected def getValue() = {
    _value
  }

@transient private lazy val _value: T = readBroadcastBlock()

可以認為TorrentBroadcast物件經過了三個主要階段的處理：構造器，序列化，反序列化

構造器

在構造TorrentBroadcast物件時，numBlocks會被初始化，此時writeBlocks會被執行。writeBlocks會執行把obj序列化，分塊，儲存進BlockManager等操作。

而_value域是lazy的，因此在TorrentBroadcast物件初始化時，_value不會初始化，readBroadcastBlock也不會執行。

序列化

當在driver端對RDD呼叫一個action時，會生成Task物件，Task物件引用到的物件會被序列化，然後對每一個task，反序列化一個Task物件。

TorrentBroadcast需要保證被廣播的物件不會隨Task一起序列化。需要注意以下兩點：

private[spark] class TorrentBroadcast[T: ClassTag](obj: T, id: Long)
  extends Broadcast[T](id) with Logging with Serializable {
　　……
}
@transient private lazy val _value: T = readBroadcastBlock()

Scala的建構函式裡的引數並不一定會成為物件的欄位，像obj這種只是用來構造物件、沒有被用於實現方法的構造器引數，不會成為TorrentBroadcast的欄位，因此不會被序列化。

而_value儘管引用了被廣播的資料，但它是@transient的，因此也不會被序列化。

反序列化

反序列化的關鍵在於，_value不會被反序列化。因此，如果某個executor沒有task使用TorrentBroadcast的value方法，被廣播的資料就不會被在這個executor端獲取。

實現這種功能的關鍵在於Scala的lazy val。

首先，考慮這個問題：lazy val可能被多個執行緒同時訪問，這會觸發lazy val的初始化，但是需要保證這個初始化的過程就執行緒安全的，即lazy val只被初始化一次，且初始化的結果對所有執行緒可見。實現這種行為，最簡單的做為是使用this做同步，但是這樣的效率會很低，而Scala實現lazy val使用了一種效率更高的方法，但不管怎麼做，lazy val比普通的val的訪問效率會降低。

舉一個Double-checked locking idiom, sweet in Scala!中的例子：

lazy val myLazyField = create();

會被編譯成：

   public volatile int bitmap$0;
   private Object myLazyField;

   public String myLazyField() {
        if((bitmap$0 & 1) == 0)
        {
            synchronized(this)
            {
                if((bitmap$0 & 1) == 0)
                {
                    myLazyField = ...
                    bitmap$0 = bitmap$0 | 1;
                }
            }
        }
        return myLazyField;
    }

即通過一個volatile變數來判斷這個lazy val是否已經初始化，通過雙重檢查加鎖來做初始化。

現在有了新的問題：

1. 預設的序列化過程是否會觸發lazy val被初始化呢？

2. 如果在TorrentBroadcast物件被序列化之前，lazy val被訪問，觸發了初始化過程，那麼被廣播的資料相關於作為TorrentBroadcast的一個field，也會被序列化。

問題1的答案是不會觸發。問題2的答案_value需要被註明是transient，就像TorrentBroadcast裡所做的一樣。

所以，在函式中如果經常使用Broadcast.value方法返回的物件時，比如在迴圈中使用它，最後先在迴圈外建立一個對這個物件的引用，以減少一些開銷。

但是，lazy val的這種執行緒安全機制對於TorrentBroadcast是浪費的。因為Broadcast變數是隨Task一起序列化的，每個執行緒有自己的Task物件，也就是執行緒間不共享Broadcast物件。實際上，為了保證同一個JVM上執行的不同task得到同樣的被廣播的物件，readBroadcastBlock方法是使用TorrentBroadcast這個class做了同步，

下面來看一下把被廣播的物件分塊儲存的過程

將廣播的物件分塊儲存

這一步是在TorrentBroadcast物件初始化時候做的。

由 

val numBlocks: Int = writeBlocks(obj)

 

觸發。下面看一下writeBlocks方法

writeBlocks

  private def writeBlocks(value: T): Int = {
    // Store a copy of the broadcast variable in the driver so that tasks run on the driver
    // do not create a duplicate copy of the broadcast variable's value.
    SparkEnv.get.blockManager.putSingle(broadcastId, value, StorageLevel.MEMORY_AND_DISK,
      tellMaster = false)
    val blocks =
      TorrentBroadcast.blockifyObject(value, blockSize, SparkEnv.get.serializer, compressionCodec)
    //blocks的型別是Array[ByteBuffer]
    blocks.zipWithIndex.foreach { case (block, i) =>
      SparkEnv.get.blockManager.putBytes(
        BroadcastBlockId(id, "piece" + i),//以BroadcastBlockId為BlockId儲存
        block,
        StorageLevel.MEMORY_AND_DISK_SER,
        tellMaster = true)
    }
    blocks.length
  }

 

正如程式碼中的註釋所說的，writeBlocks會首先把被廣播的物件用putSingle方法放在driver的BlockManager裡，這是為了當在driver執行task時，不會額外建立一個被廣播的物件的副本。若沒有這一步，在driver端執行task時，會和executor端一樣，通過Broadcast物件的value方法新建一個被廣播的物件，這就使得driver端有兩份這個物件。但實際上driver端執行task的情況並不常見。所以這裡最好根據conf判斷下是否有必要這麼做。

接下來，使用伴生物件的blockifyObject方法把物件分塊，得到的結果是一個ByteBuffer的陣列。然後把這些塊存進BlockManager, 這裡有兩點需要注意：

1. 把塊存進BlockManager時，使用的id是BroadcastBlockId(id, "piece" + i)。也就是說跟據Broadcast物件的id，以及總共的塊的數量就可以還原出所有的塊儲存時所使用的id。這也就是為什麼TorrentBroadcast要有numBlocks這個field的原因。而id欄位是Broadcast這個虛類裡的val， 所以根據TorrentBroadcast物件的欄位，即可以它所劃分的所有block的id。在從這些塊還原被broadcast的物件時，也的確是這麼做的。

2. 把劃分出的塊儲存進BlockManager時，tellMaster欄位的值為true，這就使得master可以知道哪個BlockManager儲存了這個塊，因此executor端的BlockManager最初的時候才能從driver端的BlockManager獲取這個塊。相反的是，writeBlocks第一句putSingle時，tellMaster是false，因為並不準備讓其它BlockManager獲取putSingle進去的物件。

blockifyObject

blockifyObject作的工作就是將被廣播的物件序列化，如果啟用了壓縮就進行壓縮，然後將得到的位元組流寫入到一系列位元組陣列中。

它的返回值型別為：Array[ByteBuffer], 之所有是ByteBuffer, 是為了BlockManager使用方便，因為BlockManager的putBytes方法接受ByteBuffer作為引數。

  def blockifyObject[T: ClassTag](
      obj: T,
      blockSize: Int,
      serializer: Serializer,
      compressionCodec: Option[CompressionCodec]): Array[ByteBuffer] = {
    val bos = new ByteArrayChunkOutputStream(blockSize)
    val out: OutputStream = compressionCodec.map(c => c.compressedOutputStream(bos)).getOrElse(bos)
    val ser = serializer.newInstance()
    val serOut = ser.serializeStream(out)
    serOut.writeObject[T](obj).close()
    bos.toArrays.map(ByteBuffer.wrap)
  }

它實現的關鍵在於ByteArrayChunkOutputStream, 這個類實現了Java的OutputStream介面。它的主體部分如下：

private[spark] class ByteArrayChunkOutputStream(chunkSize: Int) extends OutputStream {

  private val chunks = new ArrayBuffer[Array[Byte]]
  private var lastChunkIndex = -1
  private var position = chunkSize
  override def write(b: Int): Unit = {
    allocateNewChunkIfNeeded()
    chunks(lastChunkIndex)(position) = b.toByte
    position += 1
  }

　override def write(bytes: Array[Byte], off: Int, len: Int): Unit = { ... }
　def toArrays: Array[Array[Byte]] = { ... }
 ... 

}

即，它在內部使用一些長度等於chunkSize的陣列來儲存被寫入的位元組。

組裝還原被廣播的物件

在executor端(如果有task在driver執行的話，也可以是在driver端)需要把被切塊後的物件組裝起來，還原成被廣播的物件。這是通過對lazy val _value訪問觸發的。

 @transient private lazy val _value: T = readBroadcastBlock()

readBroadcast會首先在本地的BlockManager尋找之前存入的被廣播的物件，因此如果同一個executor中已經有task訪問過_value，那麼它就能直接取到已被放入本地BlockManager中的物件，

如果本地還沒有, 那麼就會呼叫readBlocks獲取組成這個物件的塊，然後用unblockifyObject還原這個物件，接著把它放入BlockManager，以使得同一個executor的其它task不必重複組裝還原。

 private def readBroadcastBlock(): T = Utils.tryOrIOException {
    TorrentBroadcast.synchronized {
      setConf(SparkEnv.get.conf)
      //從本地的blockManager裡讀這個被broadcast的物件，根據broadcastId
      SparkEnv.get.blockManager.getLocal(broadcastId).map(_.data.next()) match {
        case Some(x) => //本地有
          x.asInstanceOf[T]

        case None => //本地無
          logInfo("Started reading broadcast variable " + id)
          val startTimeMs = System.currentTimeMillis()
          val blocks = readBlocks()//如果本地沒有broadcastId對應的broadcast的block，就讀
          logInfo("Reading broadcast variable " + id + " took" + Utils.getUsedTimeMs(startTimeMs))

          val obj = TorrentBroadcast.unBlockifyObject[T](
            blocks, SparkEnv.get.serializer, compressionCodec)
          // Store the merged copy in BlockManager so other tasks on this executor don't
          // need to re-fetch it.
          SparkEnv.get.blockManager.putSingle( //讀了之後再放進BlockManager
            broadcastId, obj, StorageLevel.MEMORY_AND_DISK, tellMaster = false)
          obj
      }
    }
  }

這裡有一個細節是，組裝還原之後的物件被用putSingle放入BlockManager, 儲存級別為MEMORY_AND_DISK，這就意味著，在MemoryStore無法容納被廣播的物件時，同一個executor的兩個task可能會獲取兩個不同的物件(需要研究下BlockManager相關的程式碼才能確定)。如果這種情況發生，而被廣播的物件是執行緒安全的，那麼就是對記憶體的浪費。如果這種情況不發生，一個executor的所有task共享一個被廣播的物件，那麼可能會產生執行緒安全的問題。但是無論如何，使用被廣播的物件時，需要以只讀的方式，對它的修改可能會產生問題。

TorrentBroadcast是通過readBlocks獲取構成序列化後的物件的塊。

  /** Fetch torrent blocks from the driver and/or other executors. */
  private def readBlocks(): Array[ByteBuffer] = {
    //獲取到的block被存在本地的BlockManager中並且上報給driver,這樣其它的executor就可以從這個executor獲取這些block了
    val blocks = new Array[ByteBuffer](numBlocks)
    val bm = SparkEnv.get.blockManager

    //需要shuffle，避免所有executor以同樣的順序下載block，使得driver依然是瓶頸
    for (pid <- Random.shuffle(Seq.range(0, numBlocks))) {
      val pieceId = BroadcastBlockId(id, "piece" + pid)//組裝BroadcastBlockId
      logDebug(s"Reading piece $pieceId of $broadcastId")
      // 先試著從本地獲取，因為之前的嘗試可能已經獲取了一些block
      def getLocal: Option[ByteBuffer] = bm.getLocalBytes(pieceId)
      def getRemote: Option[ByteBuffer] = bm.getRemoteBytes(pieceId).map { block =>
        //如果從remote獲取了block，就把它存在本地的BlockManager
        SparkEnv.get.blockManager.putBytes(
          pieceId,
          block,
          StorageLevel.MEMORY_AND_DISK_SER,
          tellMaster = true)
        block
      }
      val block: ByteBuffer = getLocal.orElse(getRemote).getOrElse(
        throw new SparkException(s"Failed to get $pieceId of $broadcastId"))
      blocks(pid) = block
    }
    blocks
  }

readBlocks還是很簡單易懂的，只是這裡使用putBytes時，使用的儲存級別是MEMORY_AND_DISK_SER，有些奇怪，不知道為啥對於這些bytes還需要序列化。

總結

TorrentBroadcast的實現有一些巧妙的細節，但是整體的程式碼還是很簡潔，也比較容易理解。之所以有如此少的程式碼，是因為BlockManager已經提供了足夠的基礎設施。 

相關文章