### Question:
SimpleFetchedInputAllocator implements FetchedInputAllocator,
    FetchedInputCallback { @Override public synchronized FetchedInput allocate(long actualSize, long compressedSize, InputAttemptIdentifier inputAttemptIdentifier) throws IOException { if (actualSize > maxSingleShuffleLimit || this.usedMemory + actualSize > this.memoryLimit) { return new DiskFetchedInput(actualSize, compressedSize, inputAttemptIdentifier, this, conf, localDirAllocator, fileNameAllocator); } else { this.usedMemory += actualSize; LOG.info("Used memory after allocating " + actualSize + " : " + usedMemory); return new MemoryFetchedInput(actualSize, compressedSize, inputAttemptIdentifier, this); } } SimpleFetchedInputAllocator(String uniqueIdentifier, Configuration conf,
      long maxTaskAvailableMemory, long memoryAvailable); @Private static long getInitialMemoryReq(Configuration conf, long maxAvailableTaskMemory); @Override synchronized FetchedInput allocate(long actualSize, long compressedSize,
      InputAttemptIdentifier inputAttemptIdentifier); @Override synchronized void fetchComplete(FetchedInput fetchedInput); @Override synchronized void fetchFailed(FetchedInput fetchedInput); @Override synchronized void freeResources(FetchedInput fetchedInput);  }### Answer:
@Test public void testInMemAllocation() throws IOException { String localDirs = "/tmp/" + this.getClass().getName(); Configuration conf = new Configuration(); long jvmMax = Runtime.getRuntime().maxMemory(); LOG.info("jvmMax: " + jvmMax); float bufferPercent = 0.1f; conf.setFloat(TezJobConfig.TEZ_RUNTIME_SHUFFLE_INPUT_BUFFER_PERCENT, bufferPercent); conf.setFloat(TezJobConfig.TEZ_RUNTIME_SHUFFLE_MEMORY_LIMIT_PERCENT, 1.0f); conf.setStrings(TezRuntimeFrameworkConfigs.LOCAL_DIRS, localDirs); long inMemThreshold = (long) (bufferPercent * jvmMax); LOG.info("InMemThreshold: " + inMemThreshold); SimpleFetchedInputAllocator inputManager = new SimpleFetchedInputAllocator(UUID.randomUUID().toString(), conf, Runtime.getRuntime().maxMemory(), inMemThreshold); long requestSize = (long) (0.4f * inMemThreshold); long compressedSize = 1l; LOG.info("RequestSize: " + requestSize); FetchedInput fi1 = inputManager.allocate(requestSize, compressedSize, new InputAttemptIdentifier(1, 1)); assertEquals(FetchedInput.Type.MEMORY, fi1.getType()); FetchedInput fi2 = inputManager.allocate(requestSize, compressedSize, new InputAttemptIdentifier(2, 1)); assertEquals(FetchedInput.Type.MEMORY, fi2.getType()); FetchedInput fi3 = inputManager.allocate(requestSize, compressedSize, new InputAttemptIdentifier(3, 1)); assertEquals(FetchedInput.Type.DISK, fi3.getType()); fi1.abort(); fi1.free(); FetchedInput fi4 = inputManager.allocate(requestSize, compressedSize, new InputAttemptIdentifier(4, 1)); assertEquals(FetchedInput.Type.MEMORY, fi4.getType()); fi3.abort(); fi3.free(); FetchedInput fi5 = inputManager.allocate(requestSize, compressedSize, new InputAttemptIdentifier(4, 1)); assertEquals(FetchedInput.Type.DISK, fi5.getType()); }