Github user JoshRosen commented on a diff in the pull request:
https://github.com/apache/spark/pull/6227#discussion_r30663764
--- Diff:
unsafe/src/main/java/org/apache/spark/unsafe/memory/ExecutorMemoryManager.java
---
@@ -44,15 +56,56 @@ public ExecutorMemoryManager(MemoryAllocator allocator)
{
}
/**
+ * Returns true if allocations of the given size should go through the
pooling mechanism and
+ * false otherwise.
+ */
+ private boolean shouldPool(long size) {
+ // Very small allocations are less likely to benefit from pooling.
+ // At some point, we should explore supporting pooling for off-heap
memory, but for now we'll
+ // ignore that case in the interest of simplicity.
+ return size >= POOLING_THRESHOLD_BYTES && allocator instanceof
HeapMemoryAllocator;
+ }
+
+ /**
* Allocates a contiguous block of memory. Note that the allocated
memory is not guaranteed
* to be zeroed out (call `zero()` on the result if this is necessary).
*/
MemoryBlock allocate(long size) throws OutOfMemoryError {
- return allocator.allocate(size);
+ if (shouldPool(size)) {
+ synchronized (this) {
+ final LinkedList<SoftReference<MemoryBlock>> pool =
bufferPoolsBySize.get(size);
+ if (pool != null) {
+ while (!pool.isEmpty()) {
+ final SoftReference<MemoryBlock> blockReference = pool.pop();
+ final MemoryBlock memory = blockReference.get();
+ if (memory != null) {
+ assert (memory.size() == size);
+ return memory;
+ }
+ }
+ bufferPoolsBySize.remove(size);
+ }
+ }
+ return allocator.allocate(size);
+ } else {
+ return allocator.allocate(size);
+ }
}
void free(MemoryBlock memory) {
- allocator.free(memory);
+ final long size = memory.size();
+ if (shouldPool(size)) {
+ synchronized (this) {
+ LinkedList<SoftReference<MemoryBlock>> pool =
bufferPoolsBySize.get(size);
+ if (pool == null) {
+ pool = new LinkedList<SoftReference<MemoryBlock>>();
+ bufferPoolsBySize.put(size, pool);
+ }
+ pool.add(new SoftReference<MemoryBlock>(memory));
--- End diff --
Java's linked list is doubly-linked, so I don't think that this makes a
perf. difference or anything, which is why I was a little sloppy here.
---
If your project is set up for it, you can reply to this email and have your
reply appear on GitHub as well. If your project does not have this feature
enabled and wishes so, or if the feature is enabled but not working, please
contact infrastructure at [email protected] or file a JIRA ticket
with INFRA.
---
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
