0001
0002
0003
0004
0005
0006
0007
0008
0009
0010
0011
0012
0013
0014
0015
0016
0017
0018 package org.apache.spark.memory;
0019
0020 import com.google.common.annotations.VisibleForTesting;
0021
0022 import org.apache.spark.unsafe.memory.MemoryBlock;
0023
0024 import java.io.IOException;
0025
0026 public class TestMemoryConsumer extends MemoryConsumer {
0027 public TestMemoryConsumer(TaskMemoryManager memoryManager, MemoryMode mode) {
0028 super(memoryManager, 1024L, mode);
0029 }
0030 public TestMemoryConsumer(TaskMemoryManager memoryManager) {
0031 this(memoryManager, MemoryMode.ON_HEAP);
0032 }
0033
0034 @Override
0035 public long spill(long size, MemoryConsumer trigger) throws IOException {
0036 long used = getUsed();
0037 free(used);
0038 return used;
0039 }
0040
0041 public void use(long size) {
0042 long got = taskMemoryManager.acquireExecutionMemory(size, this);
0043 used += got;
0044 }
0045
0046 public void free(long size) {
0047 used -= size;
0048 taskMemoryManager.releaseExecutionMemory(size, this);
0049 }
0050
0051 @VisibleForTesting
0052 public void freePage(MemoryBlock page) {
0053 used -= page.size();
0054 taskMemoryManager.freePage(page, this);
0055 }
0056 }
0057
0058