spark 2.1 MemoryPool, StoragememoryPool

MemoryPool

/** * Manages bookkeeping for an adjustable-sized region of memory. This class is internal to * the [[MemoryManager]]. See subclasses for more details. * * @param lock a [[MemoryManager]] instance, used for synchronization. We purposely erase the type *             to `Object` to avoid programming errors, since this object should only be used for *             synchronization purposes. */private[memory] abstract class MemoryPool(lock: Object) {  @GuardedBy("lock")  private[this] var _poolSize: Long = 0  /**   * Returns the current size of the pool, in bytes.   */  final def poolSize: Long = lock.synchronized {    _poolSize  }  /**   * Returns the amount of free memory in the pool, in bytes.   */  final def memoryFree: Long = lock.synchronized {    _poolSize - memoryUsed  }  /**   * Expands the pool by `delta` bytes.   */  final def incrementPoolSize(delta: Long): Unit = lock.synchronized {    require(delta >= 0)    _poolSize += delta  }  /**   * Shrinks the pool by `delta` bytes.   */  final def decrementPoolSize(delta: Long): Unit = lock.synchronized {    require(delta >= 0)    require(delta <= _poolSize)    require(_poolSize - delta >= memoryUsed)    _poolSize -= delta  }  /**   * Returns the amount of used memory in this pool (in bytes).   */  def memoryUsed: Long}

MemoryMode

@Privatepublic enum MemoryMode {  ON_HEAP,  OFF_HEAP}

StorageMemoryPool

/** * Performs bookkeeping for managing an adjustable-size pool of memory that is used for storage * (caching). * * @param lock a [[MemoryManager]] instance to synchronize on * @param memoryMode the type of memory tracked by this pool (on- or off-heap) */private[memory] class StorageMemoryPool(    lock: Object,    memoryMode: MemoryMode  ) extends MemoryPool(lock) with Logging {  private[this] val poolName: String = memoryMode match {    case MemoryMode.ON_HEAP => "on-heap storage"    case MemoryMode.OFF_HEAP => "off-heap storage"  }  @GuardedBy("lock")  private[this] var _memoryUsed: Long = 0L  override def memoryUsed: Long = lock.synchronized {    _memoryUsed  }  private var _memoryStore: MemoryStore = _  def memoryStore: MemoryStore = {    if (_memoryStore == null) {      throw new IllegalStateException("memory store not initialized yet")    }    _memoryStore  }  /**   * Set the [[MemoryStore]] used by this manager to evict cached blocks.   * This must be set after construction due to initialization ordering constraints.   */  final def setMemoryStore(store: MemoryStore): Unit = {    _memoryStore = store  }

StoragememoryPool.acquireMemory

  /**   * Acquire N bytes of memory to cache the given block, evicting existing ones if necessary.   *   * @return whether all N bytes were successfully granted.   */  def acquireMemory(blockId: BlockId, numBytes: Long): Boolean = lock.synchronized {    val numBytesToFree = math.max(0, numBytes - memoryFree)    acquireMemory(blockId, numBytes, numBytesToFree)  }

  /**   * Acquire N bytes of storage memory for the given block, evicting existing ones if necessary.   *   * @param blockId the ID of the block we are acquiring storage memory for   * @param numBytesToAcquire the size of this block   * @param numBytesToFree the amount of space to be freed through evicting blocks   * @return whether all N bytes were successfully granted.   */  def acquireMemory(      blockId: BlockId,      numBytesToAcquire: Long,      numBytesToFree: Long): Boolean = lock.synchronized {    assert(numBytesToAcquire >= 0)    assert(numBytesToFree >= 0)    assert(memoryUsed <= poolSize)    if (numBytesToFree > 0) {      memoryStore.evictBlocksToFreeSpace(Some(blockId), numBytesToFree, memoryMode)    }    // NOTE: If the memory store evicts blocks, then those evictions will synchronously call    // back into this StorageMemoryPool in order to free memory. Therefore, these variables    // should have been updated.    val enoughMemory = numBytesToAcquire <= memoryFree    if (enoughMemory) {      _memoryUsed += numBytesToAcquire    }    enoughMemory  }

acquireMemory

/**   * Acquire N bytes of storage memory for the given block, evicting existing ones if necessary.   *   * @param blockId the ID of the block we are acquiring storage memory for   * @param numBytesToAcquire the size of this block   * @param numBytesToFree the amount of space to be freed through evicting blocks   * @return whether all N bytes were successfully granted.   */  def acquireMemory(      blockId: BlockId,      numBytesToAcquire: Long,      numBytesToFree: Long): Boolean = lock.synchronized {    assert(numBytesToAcquire >= 0)    assert(numBytesToFree >= 0)    assert(memoryUsed <= poolSize)    if (numBytesToFree > 0) {      memoryStore.evictBlocksToFreeSpace(Some(blockId), numBytesToFree, memoryMode)    }    // NOTE: If the memory store evicts blocks, then those evictions will synchronously call    // back into this StorageMemoryPool in order to free memory. Therefore, these variables    // should have been updated.    val enoughMemory = numBytesToAcquire <= memoryFree    if (enoughMemory) {      _memoryUsed += numBytesToAcquire    }    enoughMemory  }

  def releaseMemory(size: Long): Unit = lock.synchronized {    if (size > _memoryUsed) {      logWarning(s"Attempted to release $size bytes of storage " +        s"memory when we only have ${_memoryUsed} bytes")      _memoryUsed = 0    } else {      _memoryUsed -= size    }  }  def releaseAllMemory(): Unit = lock.synchronized {    _memoryUsed = 0  }

/**   * Free space to shrink the size of this storage memory pool by `spaceToFree` bytes.   * Note: this method doesn't actually reduce the pool size but relies on the caller to do so.   *   * @return number of bytes to be removed from the pool's capacity.   */  def freeSpaceToShrinkPool(spaceToFree: Long): Long = lock.synchronized {    val spaceFreedByReleasingUnusedMemory = math.min(spaceToFree, memoryFree)    val remainingSpaceToFree = spaceToFree - spaceFreedByReleasingUnusedMemory    if (remainingSpaceToFree > 0) {      // If reclaiming free memory did not adequately shrink the pool, begin evicting blocks:      val spaceFreedByEviction =        memoryStore.evictBlocksToFreeSpace(None, remainingSpaceToFree, memoryMode)      // When a block is released, BlockManager.dropFromMemory() calls releaseMemory(), so we do      // not need to decrement _memoryUsed here. However, we do need to decrement the pool size.      spaceFreedByReleasingUnusedMemory + spaceFreedByEviction    } else {      spaceFreedByReleasingUnusedMemory    }  }