TaskMemoryManager

jaceklaskowski · jaceklaskowski · commit 1f89b0297fec · 2021-08-20T23:02:08.000+02:00
diff --git a/docs/memory/MemoryAllocator.md b/docs/memory/MemoryAllocator.md
@@ -0,0 +1,3 @@
+# MemoryAllocator
+
+`MemoryAllocator` is...FIXME
diff --git a/docs/memory/TaskMemoryManager.md b/docs/memory/TaskMemoryManager.md
@@ -2,6 +2,13 @@
 
 `TaskMemoryManager` manages the memory allocated to a [single task](#taskAttemptId) (using [MemoryManager](#memoryManager)).
 
+`TaskMemoryManager` assumes that:
+
+1. <span id="PAGE_NUMBER_BITS"> The number of bits to address pages is `13`
+1. <span id="OFFSET_BITS"> The number of bits to encode offsets in pages is `51` (64 bits - [13 bits](#PAGE_NUMBER_BITS))
+1. <span id="PAGE_TABLE_SIZE"> Number of pages in the [page table](#pageTable) and [to be allocated](#allocatedPages) is `8192` (`1 <<` [13](#PAGE_NUMBER_BITS))
+1. <span id="MAXIMUM_PAGE_SIZE_BYTES"> The maximum page size is `15GB` (`((1L << 31) - 1) * 8L`)
+
 ## Creating Instance
 
 `TaskMemoryManager` takes the following to be created:
@@ -28,6 +35,21 @@
 * [freePage](#freePage)
 * [getMemoryConsumptionForThisTask](#getMemoryConsumptionForThisTask)
 
+## <span id="pageTable"> Page Table (MemoryBlocks)
+
+`TaskMemoryManager` uses an array of `MemoryBlock`s (to mimic an operating system's page table).
+
+The page table uses 13 bits for addressing pages.
+
+A page is "stored" in [allocatePage](#allocatePage) and "removed" in [freePage](#freePage).
+
+All pages are released (_removed_) in [cleanUpAllAllocatedMemory](#cleanUpAllAllocatedMemory).
+
+`TaskMemoryManager` uses the page table when requested to:
+
+* [getPage](#getPage)
+* [getOffsetInPage](#getOffsetInPage)
+
 ## <span id="consumers"> Spillable Memory Consumers
 
 ```java
@@ -42,6 +64,44 @@ HashSet<MemoryConsumer> consumers
 
 Memory consumers are used to report memory usage when `TaskMemoryManager` is requested to [show memory usage](#showMemoryUsage).
 
+## <span id="acquiredButNotUsed"> Memory Acquired But Not Used
+
+`TaskMemoryManager` tracks the size of memory [allocated](#allocatePage) but not used (by any of the [MemoryConsumer](#consumers)s due to a `OutOfMemoryError` upon trying to use it).
+
+`TaskMemoryManager` releases the memory when [cleaning up all the allocated memory](#cleanUpAllAllocatedMemory).
+
+## <span id="allocatedPages"> Allocated Pages
+
+```java
+BitSet allocatedPages
+```
+
+`TaskMemoryManager` uses a `BitSet` ([Java]({{ java.api }}/java.base.java/util/BitSet.html)) to track [allocated pages](#allocatePage).
+
+The size is exactly the number of entries in the [page table](#pageTable) ([8192](#PAGE_TABLE_SIZE)).
+
+## <span id="tungstenMemoryMode"><span id="getTungstenMemoryMode"> MemoryMode
+
+`TaskMemoryManager` can be in `ON_HEAP` or `OFF_HEAP` mode (to avoid extra work for off-heap and hoping that the JIT handles branching well).
+
+`TaskMemoryManager` is given the `MemoryMode` matching the [MemoryMode](MemoryManager.md#tungstenMemoryMode) (of the given [MemoryManager](#memoryManager)) when [created](#creating-instance).
+
+`TaskMemoryManager` uses the `MemoryMode` to match to for the following:
+
+* [allocatePage](#allocatePage)
+* [cleanUpAllAllocatedMemory](#cleanUpAllAllocatedMemory)
+
+For `OFF_HEAP` mode, `TaskMemoryManager` has to change offset while [encodePageNumberAndOffset](#encodePageNumberAndOffset) and [getOffsetInPage](#getOffsetInPage).
+
+For `OFF_HEAP` mode, `TaskMemoryManager` returns no [page](#getPage).
+
+The `MemoryMode` is used when:
+
+* `ShuffleExternalSorter` is [created](../shuffle/ShuffleExternalSorter.md)
+* `BytesToBytesMap` is [created](BytesToBytesMap.md)
+* `UnsafeExternalSorter` is [created](UnsafeExternalSorter.md)
+* `Spillable` is requested to [spill](../shuffle/Spillable.md#spill) (only when in `ON_HEAP` mode)
+
 ## <span id="acquireExecutionMemory"> Acquiring Execution Memory
 
 ```java
@@ -162,181 +222,151 @@ Acquired by [consumer]: [memUsage]
 
 * `MemoryConsumer` is requested to [throw an OutOfMemoryError](MemoryConsumer.md#throwOom)
 
-## Logging
-
-Enable `ALL` logging level for `org.apache.spark.memory.TaskMemoryManager` logger to see what happens inside.
-
-Add the following line to `conf/log4j.properties`:
+## <span id="cleanUpAllAllocatedMemory"> Cleaning Up All Allocated Memory
 
-```text
-log4j.logger.org.apache.spark.memory.TaskMemoryManager=ALL
+```java
+long cleanUpAllAllocatedMemory()
 ```
 
-Refer to [Logging](../spark-logging.md).
+The `consumers` collection is then cleared.
 
-## Review Me
+`cleanUpAllAllocatedMemory` finds all the registered [MemoryConsumer](MemoryConsumer.md)s (in the [consumers](#consumers) registry) that still keep [some memory used](MemoryConsumer.md#getUsed) and, for every such consumer, prints out the following DEBUG message to the logs:
 
-TaskMemoryManager assumes that:
+```text
+unreleased [getUsed] memory from [consumer]
+```
 
-* The number of bits to address pages (aka `PAGE_NUMBER_BITS`) is `13`
-* The number of bits to encode offsets in data pages (aka `OFFSET_BITS`) is `51` (i.e. 64 bits - `PAGE_NUMBER_BITS`)
-* The number of entries in the <<pageTable, page table>> and <<allocatedPages, allocated pages>> (aka `PAGE_TABLE_SIZE`) is `8192` (i.e. 1 << `PAGE_NUMBER_BITS`)
-* The maximum page size (aka `MAXIMUM_PAGE_SIZE_BYTES`) is `15GB` (i.e. `((1L << 31) - 1) * 8L`)
+`cleanUpAllAllocatedMemory` removes all the [consumers](#consumers).
 
-== [[cleanUpAllAllocatedMemory]] Cleaning Up All Allocated Memory
+---
 
-[source, java]
-----
-long cleanUpAllAllocatedMemory()
-----
+For every `MemoryBlock` in the [pageTable](#pageTable), `cleanUpAllAllocatedMemory` prints out the following DEBUG message to the logs:
 
-`cleanUpAllAllocatedMemory` clears <<pageTable, page table>>.
+```text
+unreleased page: [page] in task [taskAttemptId]
+```
 
-CAUTION: FIXME
+`cleanUpAllAllocatedMemory` marks the pages to be freed (`FREED_IN_TMM_PAGE_NUMBER`) and requests the [MemoryManager](#memoryManager) for the [tungstenMemoryAllocator](MemoryManager.md#tungstenMemoryAllocator) to [free up the MemoryBlock](MemoryAllocator.md#free).
 
-All recorded <<consumers, consumers>> are queried for the size of used memory. If the memory used is greater than 0, the following WARN message is printed out to the logs:
+`cleanUpAllAllocatedMemory` clears the [pageTable](#pageTable) registry (by assigning `null` values).
 
-```
-WARN TaskMemoryManager: leak [bytes] memory from [consumer]
-```
+---
 
-The `consumers` collection is then cleared.
+`cleanUpAllAllocatedMemory` requests the [MemoryManager](#memoryManager) to [release execution memory](MemoryManager.md#releaseExecutionMemory) that is not used by any consumer (with the [acquiredButNotUsed](#acquiredButNotUsed) and the [tungstenMemoryMode](#tungstenMemoryMode)).
 
-MemoryManager.md#releaseExecutionMemory[MemoryManager.releaseExecutionMemory] is executed to release the memory that is not used by any consumer.
+In the end, `cleanUpAllAllocatedMemory` requests the [MemoryManager](#memoryManager) to [release all execution memory for the task](MemoryManager.md#releaseAllExecutionMemoryForTask).
 
-Before `cleanUpAllAllocatedMemory` returns, it calls MemoryManager.md#releaseAllExecutionMemoryForTask[MemoryManager.releaseAllExecutionMemoryForTask] that in turn becomes the return value.
+---
 
-CAUTION: FIXME Image with the interactions to `MemoryManager`.
+`cleanUpAllAllocatedMemory` is used when:
 
-NOTE: `cleanUpAllAllocatedMemory` is used exclusively when `TaskRunner` is requested to executor:TaskRunner.md#run[run] (and cleans up after itself).
+* `TaskRunner` is requested to [run a task](../executor/TaskRunner.md#run) (and the task has finished successfully)
 
-== [[allocatePage]] Allocating Memory Block for Tungsten Consumers
+## <span id="allocatePage"> Allocating Memory Page
 
-[source, java]
-----
+```java
 MemoryBlock allocatePage(
   long size,
   MemoryConsumer consumer)
-----
-
-NOTE: It only handles *Tungsten Consumers*, i.e. MemoryConsumer.md[MemoryConsumers] in `tungstenMemoryMode` mode.
-
-`allocatePage` allocates a block of memory (aka _page_) smaller than `MAXIMUM_PAGE_SIZE_BYTES` maximum size.
-
-It checks `size` against the internal `MAXIMUM_PAGE_SIZE_BYTES` maximum size. If it is greater than the maximum size, the following `IllegalArgumentException` is thrown:
-
-```
-Cannot allocate a page with more than [MAXIMUM_PAGE_SIZE_BYTES] bytes
 ```
 
-It then <<acquireExecutionMemory, acquires execution memory>> (for the input `size` and `consumer`).
-
-It finishes by returning `null` when no execution memory could be acquired.
-
-With the execution memory acquired, it finds the smallest unallocated page index and records the page number (using <<allocatedPages, allocatedPages>> registry).
+`allocatePage` allocates a block of memory (_page_) that is:
 
-If the index is `PAGE_TABLE_SIZE` or higher, <<releaseExecutionMemory, releaseExecutionMemory(acquired, consumer)>> is called and then the following `IllegalStateException` is thrown:
+1. Below [MAXIMUM_PAGE_SIZE_BYTES](#MAXIMUM_PAGE_SIZE_BYTES) maximum size
+1. For [MemoryConsumer](MemoryConsumer.md)s with the same [MemoryMode](MemoryConsumer.md#getMode) as the [TaskMemoryManager](#tungstenMemoryMode)
 
-```
-Have already allocated a maximum of [PAGE_TABLE_SIZE] pages
-```
+`allocatePage` [acquireExecutionMemory](#acquireExecutionMemory) (for the `size` and the [MemoryConsumer](MemoryConsumer.md)). `allocatePage` returns immediately (with `null`) when this allocation ended up with `0` or less bytes.
 
-It then attempts to allocate a `MemoryBlock` from `Tungsten MemoryAllocator` (calling `memoryManager.tungstenMemoryAllocator().allocate(acquired)`).
+`allocatePage` allocates the first clear bit in the [allocatedPages](#allocatedPages) (unless the whole page table is taken and `allocatePage` throws an `IllegalStateException`).
 
-CAUTION: FIXME What is `MemoryAllocator`?
+`allocatePage` requests the [MemoryManager](#memoryManager) for the [tungstenMemoryAllocator](MemoryManager.md#tungstenMemoryAllocator) that is requested to [allocate the acquired memory](MemoryAllocator.md#allocate).
 
-When successful, `MemoryBlock` gets assigned `pageNumber` and it gets added to the internal <<pageTable, pageTable>> registry.
+`allocatePage` registers the page in the [pageTable](#pageTable).
 
-You should see the following TRACE message in the logs:
+In the end, `allocatePage` prints out the following TRACE message to the logs and returns the `MemoryBlock` allocated.
 
-```
-TRACE Allocate page number [pageNumber] ([acquired] bytes)
+```text
+Allocate page number [pageNumber] ([acquired] bytes)
 ```
 
-The `page` is returned.
+### <span id="allocatePage-usage"> Usage
 
-If a `OutOfMemoryError` is thrown when allocating a `MemoryBlock` page, the following WARN message is printed out to the logs:
+`allocatePage` is used when:
 
-```
-WARN Failed to allocate a page ([acquired] bytes), try again.
-```
-
-And `acquiredButNotUsed` gets `acquired` memory space with the `pageNumber` cleared in <<allocatedPages, allocatedPages>> (i.e. the index for `pageNumber` gets `false`).
-
-CAUTION: FIXME Why is the code tracking `acquiredButNotUsed`?
+* `MemoryConsumer` is requested to allocate an [array](MemoryConsumer.md#allocateArray) and a [page](MemoryConsumer.md#allocatePage)
 
-Another <<allocatePage, allocatePage>> attempt is recursively tried.
+### <span id="allocatePage-TooLargePageException"> TooLargePageException
 
-CAUTION: FIXME Why is there a hope for being able to allocate a page?
+For sizes larger than the [MAXIMUM_PAGE_SIZE_BYTES](#MAXIMUM_PAGE_SIZE_BYTES) `allocatePage` throws a `TooLargePageException`.
 
-== [[getMemoryConsumptionForThisTask]] `getMemoryConsumptionForThisTask` Method
+### <span id="allocatePage-OutOfMemoryError"> OutOfMemoryError
 
-[source, java]
-----
-long getMemoryConsumptionForThisTask()
-----
+Requesting the [tungstenMemoryAllocator](MemoryManager.md#tungstenMemoryAllocator) to [allocate the acquired memory](MemoryAllocator.md#allocate) may throw an `OutOfMemoryError`. If so, `allocatePage` prints out the following WARN message to the logs:
 
-`getMemoryConsumptionForThisTask`...FIXME
+```text
+Failed to allocate a page ([acquired] bytes), try again.
+```
 
-NOTE: `getMemoryConsumptionForThisTask` is used exclusively in Spark tests.
+`allocatePage` adds the acquired memory to the [acquiredButNotUsed](#acquiredButNotUsed) and removes the page from the [allocatedPages](#allocatedPages) (by clearing the bit).
 
-== [[freePage]] Freeing Memory Page -- `freePage` Method
+In the end, `allocatePage` tries to [allocate the page](#allocatePage) again (recursively).
 
-[source, java]
-----
-void freePage(MemoryBlock page, MemoryConsumer consumer)
-----
+## <span id="freePage"> Freeing Memory Page
 
-`pageSizeBytes` simply requests the <<memoryManager, MemoryManager>> for MemoryManager.md#pageSizeBytes[pageSizeBytes].
+```java
+void freePage(
+  MemoryBlock page,
+  MemoryConsumer consumer)
+```
 
-NOTE: `pageSizeBytes` is used when `MemoryConsumer` is requested to MemoryConsumer.md#freePage[freePage] and MemoryConsumer.md#throwOom[throwOom].
+`pageSizeBytes` requests the [MemoryManager](#memoryManager) for [pageSizeBytes](MemoryManager.md#pageSizeBytes).
 
-== [[getPage]] Getting Page -- `getPage` Method
+`pageSizeBytes` is used when:
 
-[source, java]
-----
-Object getPage(long pagePlusOffsetAddress)
-----
+* `MemoryConsumer` is requested to [freePage](MemoryConsumer.md#freePage) and [throwOom](MemoryConsumer.md#throwOom)
 
-`getPage`...FIXME
+## <span id="getPage"> Getting Page
 
-NOTE: `getPage` is used when...FIXME
+```java
+Object getPage(
+  long pagePlusOffsetAddress)
+```
 
-== [[getOffsetInPage]] Getting Page Offset -- `getOffsetInPage` Method
+`getPage` handles the `ON_HEAP` mode of the [tungstenMemoryMode](#tungstenMemoryMode) only.
 
-[source, java]
-----
-long getOffsetInPage(long pagePlusOffsetAddress)
-----
+`getPage` looks up the page (by the given address) in the [page table](#pageTable) and requests it for the base object.
 
-`getPage`...FIXME
+`getPage` is used when:
 
-NOTE: `getPage` is used when...FIXME
+* `ShuffleExternalSorter` is requested to [writeSortedFile](../shuffle/ShuffleExternalSorter.md#writeSortedFile)
+* `Location` (of [BytesToBytesMap](BytesToBytesMap.md)) is requested to `updateAddressesAndSizes`
+* `SortComparator` (of [UnsafeInMemorySorter](UnsafeInMemorySorter.md)) is requested to `compare` two record pointers
+* `SortedIterator` (of [UnsafeInMemorySorter](UnsafeInMemorySorter.md)) is requested to `loadNext` record
 
-== [[internal-properties]] Internal Properties
+## <span id="getOffsetInPage"> getOffsetInPage
 
-[cols="30m,70",options="header",width="100%"]
-|===
-| Name
-| Description
+```java
+long getOffsetInPage(
+  long pagePlusOffsetAddress)
+```
 
-| acquiredButNotUsed
-| [[acquiredButNotUsed]] The size of memory allocated but not used.
+`getOffsetInPage` gives the offset associated with the given `pagePlusOffsetAddress` (encoded by `encodePageNumberAndOffset`).
 
-| allocatedPages
-| [[allocatedPages]] Collection of flags (`true` or `false` values) of size `PAGE_TABLE_SIZE` with all bits initially disabled (i.e. `false`).
+`getOffsetInPage` is used when:
 
-TIP: `allocatedPages` is https://docs.oracle.com/javase/8/docs/api/java/util/BitSet.html[java.util.BitSet].
+* `ShuffleExternalSorter` is requested to [writeSortedFile](../shuffle/ShuffleExternalSorter.md#writeSortedFile)
+* `Location` (of [BytesToBytesMap](BytesToBytesMap.md)) is requested to `updateAddressesAndSizes`
+* `SortComparator` (of [UnsafeInMemorySorter](UnsafeInMemorySorter.md)) is requested to `compare` two record pointers
+* `SortedIterator` (of [UnsafeInMemorySorter](UnsafeInMemorySorter.md)) is requested to `loadNext` record
 
-When <<allocatePage, allocatePage>> is called, it will record the page in the registry by setting the bit at the specified index (that corresponds to the allocated page) to `true`.
+## Logging
 
-| pageTable
-| [[pageTable]] The array of size `PAGE_TABLE_SIZE` with indices being `MemoryBlock` objects.
+Enable `ALL` logging level for `org.apache.spark.memory.TaskMemoryManager` logger to see what happens inside.
 
-When <<allocatePage, allocating a `MemoryBlock` page for Tungsten consumers>>, the index corresponds to `pageNumber` that points to the `MemoryBlock` page allocated.
+Add the following line to `conf/log4j.properties`:
 
-| tungstenMemoryMode
-| [[tungstenMemoryMode]] `MemoryMode` (i.e. `OFF_HEAP` or `ON_HEAP`)
+```text
+log4j.logger.org.apache.spark.memory.TaskMemoryManager=ALL
+```
 
-Set to the MemoryManager.md#tungstenMemoryMode[tungstenMemoryMode] of the <<memoryManager, MemoryManager>> while TaskMemoryManager is <<creating-instance, created>>
-|===
+Refer to [Logging](../spark-logging.md).
diff --git a/docs/storage/MemoryStore.md b/docs/storage/MemoryStore.md
@@ -301,13 +301,13 @@ remove(
   blockId: BlockId): Boolean
 ```
 
-`remove` returns `true` when the given block ([BlockId](BlockId.md)) was (found and) removed from the [entries](#entries) registry successfully and the [memory released](MemoryManager.md#releaseStorageMemory) (from the [MemoryManager](#memoryManager)).
+`remove` returns `true` when the given block ([BlockId](BlockId.md)) was (found and) removed from the [entries](#entries) registry successfully and the [memory released](../memory/MemoryManager.md#releaseStorageMemory) (from the [MemoryManager](#memoryManager)).
 
 ---
 
 `remove` removes (_drops_) the block ([BlockId](BlockId.md)) from the [entries](#entries) registry.
 
-If found and removed, `remove` requests the [MemoryManager](#memoryManager) to [releaseStorageMemory](MemoryManager.md#releaseStorageMemory) and prints out the following DEBUG message to the logs (with the [maxMemory](#maxMemory) and [blocksMemoryUsed](#blocksMemoryUsed)):
+If found and removed, `remove` requests the [MemoryManager](#memoryManager) to [releaseStorageMemory](../memory/MemoryManager.md#releaseStorageMemory) and prints out the following DEBUG message to the logs (with the [maxMemory](#maxMemory) and [blocksMemoryUsed](#blocksMemoryUsed)):
 
 ```text
 Block [blockId] of size [size] dropped from memory (free [memory])
diff --git a/mkdocs.yml b/mkdocs.yml
@@ -282,6 +282,7 @@ nav:
       - memory/index.md
       - BytesToBytesMap: memory/BytesToBytesMap.md
       - ExecutionMemoryPool: memory/ExecutionMemoryPool.md
+      - MemoryAllocator: memory/MemoryAllocator.md
       - MemoryConsumer: memory/MemoryConsumer.md
       - MemoryManager: memory/MemoryManager.md
       - MemoryPool: memory/MemoryPool.md

Original file line number	Diff line number	Diff line change
`@@ -0,0 +1,3 @@`
	`1`	`+# MemoryAllocator`
	`2`	`+`
	`3`	+`MemoryAllocator` is...FIXME