Java Code Examples for org.apache.flink.runtime.io.network.buffer.NetworkBufferPool#createBufferPool()

@Before
public void setUp() throws Exception {
	networkBufferPool = new NetworkBufferPool(100, 8192);
	testBufferPool = networkBufferPool.createBufferPool(1, Integer.MAX_VALUE);

	final Configuration configuration = new Configuration();
	configuration.setInteger(WebOptions.BACKPRESSURE_NUM_SAMPLES, BACKPRESSURE_NUM_SAMPLES);

	testingMiniCluster = new TestingMiniCluster(new TestingMiniClusterConfiguration.Builder()
		.setNumTaskManagers(JOB_PARALLELISM)
		.setConfiguration(configuration)
		.build());
	testingMiniCluster.start();

	dispatcherGateway = testingMiniCluster.getDispatcherGatewayFuture().get(TIMEOUT_SECONDS, TimeUnit.SECONDS);
}
 

@Test
public void testDoublePartitionRequest() throws Exception {
	final PartitionRequestClientHandler handler = new PartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = new PartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, client);

	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		final int numExclusiveBuffers = 2;
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);
		inputChannel.requestSubpartition(0);

		// The input channel should only send one partition request
		assertTrue(channel.isWritable());
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(numExclusiveBuffers, ((PartitionRequest) readFromOutbound).credit);

		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

/**
 * Verifies that {@link RemoteInputChannel#onBuffer(Buffer, int, int)} is called when a
 * {@link BufferResponse} is received.
 */
@Test
public void testReceiveBuffer() throws Exception {
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32, 2);
	final SingleInputGate inputGate = createSingleInputGate(1, networkBufferPool);
	final RemoteInputChannel inputChannel = InputChannelBuilder.newBuilder().buildRemoteChannel(inputGate);
	try {
		inputGate.setInputChannels(inputChannel);
		final BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();

		final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
		handler.addInputChannel(inputChannel);

		final int backlog = 2;
		final BufferResponse bufferResponse = createBufferResponse(
			TestBufferFactory.createBuffer(32),
			0,
			inputChannel.getInputChannelId(),
			backlog,
			new NetworkBufferAllocator(handler));
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
		assertEquals(2, inputChannel.getSenderBacklog());
	} finally {
		releaseResource(inputGate, networkBufferPool);
	}
}
 

/**
 * Tests that the RecordWriter is available iif the respective LocalBufferPool has at-least one available buffer.
 */
@Test
public void testIsAvailableOrNot() throws Exception {
	// setup
	final NetworkBufferPool globalPool = new NetworkBufferPool(10, 128, 2);
	final BufferPool localPool = globalPool.createBufferPool(1, 1, null, 1, Integer.MAX_VALUE);
	final ResultPartitionWriter resultPartition = new ResultPartitionBuilder()
		.setBufferPoolFactory(p -> localPool)
		.build();
	resultPartition.setup();
	final ResultPartitionWriter partitionWrapper = new ConsumableNotifyingResultPartitionWriterDecorator(
		new NoOpTaskActions(),
		new JobID(),
		resultPartition,
		new NoOpResultPartitionConsumableNotifier());
	final RecordWriter recordWriter = createRecordWriter(partitionWrapper);

	try {
		// record writer is available because of initial available global pool
		assertTrue(recordWriter.getAvailableFuture().isDone());

		// request one buffer from the local pool to make it unavailable afterwards
		final BufferBuilder bufferBuilder = resultPartition.getBufferBuilder(0);
		assertNotNull(bufferBuilder);
		assertFalse(recordWriter.getAvailableFuture().isDone());

		// recycle the buffer to make the local pool available again
		final Buffer buffer = BufferBuilderTestUtils.buildSingleBuffer(bufferBuilder);
		buffer.recycleBuffer();
		assertTrue(recordWriter.getAvailableFuture().isDone());
		assertEquals(recordWriter.AVAILABLE, recordWriter.getAvailableFuture());

	} finally {
		localPool.lazyDestroy();
		globalPool.destroy();
	}
}
 

@Test
public void testResumeConsumption() throws Exception {
	final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = createPartitionRequestClient(channel, handler);

	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32, 2);
	final SingleInputGate inputGate = createSingleInputGate(1, networkBufferPool);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, client);

	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		inputChannel.resumeConsumption();
		channel.runPendingTasks();
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));

		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(ResumeConsumption.class));
		assertEquals(inputChannel.getInputChannelId(), ((ResumeConsumption) readFromOutbound).receiverId);

		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.close();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

@Test
public void testDoublePartitionRequest() throws Exception {
	final PartitionRequestClientHandler handler = new PartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = new NettyPartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	final int numExclusiveBuffers = 2;
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32, numExclusiveBuffers);
	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, client, networkBufferPool);

	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		// The input channel should only send one partition request
		assertTrue(channel.isWritable());
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(numExclusiveBuffers, ((PartitionRequest) readFromOutbound).credit);

		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.close();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

private RecordWriter createRecordWriter(NetworkBufferPool globalPool) throws Exception {
	final BufferPool localPool = globalPool.createBufferPool(1, 1, null, 1, Integer.MAX_VALUE);
	final ResultPartitionWriter partition = new ResultPartitionBuilder()
		.setBufferPoolFactory(p -> localPool)
		.build();
	partition.setup();

	return new RecordWriterBuilder().build(partition);
}
 

/**
 * Verifies that {@link RemoteInputChannel#onBuffer(Buffer, int, int)} is called when a
 * {@link BufferResponse} is received.
 */
@Test
public void testReceiveBuffer() throws Exception {
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		final int numExclusiveBuffers = 2;
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);

		final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
		handler.addInputChannel(inputChannel);

		final int backlog = 2;
		final BufferResponse bufferResponse = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel.getInputChannelId(), backlog);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
		assertEquals(2, inputChannel.getSenderBacklog());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

/**
 * Verifies that {@link RemoteInputChannel#onBuffer(Buffer, int, int)} is called when a
 * {@link BufferResponse} is received.
 */
@Test
public void testReceiveBuffer() throws Exception {
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		final int numExclusiveBuffers = 2;
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);

		final PartitionRequestClientHandler handler = new PartitionRequestClientHandler();
		handler.addInputChannel(inputChannel);

		final int backlog = 2;
		final BufferResponse bufferResponse = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel.getInputChannelId(), backlog);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

/**
 * Tests to verify that there is no race condition with two things running in parallel:
 * requesting floating buffers on sender backlog and some other thread recycling
 * floating or exclusive buffers.
 */
@Test
public void testConcurrentOnSenderBacklogAndRecycle() throws Exception {
	// Setup
	final int numExclusiveSegments = 120;
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(248, 32, numExclusiveSegments);
	final int numFloatingBuffers = 128;
	final int backlog = 128;

	final ExecutorService executor = Executors.newFixedThreadPool(3);

	final SingleInputGate inputGate = createSingleInputGate(1, networkBufferPool);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	inputGate.setInputChannels(inputChannel);
	Throwable thrown = null;
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		final Callable<Void> requestBufferTask = new Callable<Void>() {
			@Override
			public Void call() throws Exception {
				for (int j = 1; j <= backlog; j++) {
					inputChannel.onSenderBacklog(j);
				}

				return null;
			}
		};

		// Submit tasks and wait to finish
		submitTasksAndWaitForResults(executor, new Callable[]{
			recycleBufferTask(inputChannel, bufferPool, numExclusiveSegments, numFloatingBuffers),
			requestBufferTask});

		assertEquals("There should be " + inputChannel.getNumberOfRequiredBuffers() + " buffers available in channel.",
			inputChannel.getNumberOfRequiredBuffers(), inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be no buffers available in local pool.",
			0, bufferPool.getNumberOfAvailableMemorySegments());
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, executor, null, thrown, inputChannel);
	}
}
 

/**
 * Tests to verify that there is no race condition with two things running in parallel:
 * requesting floating buffers on sender backlog and some other thread releasing
 * the input channel.
 */
@Test
public void testConcurrentOnSenderBacklogAndRelease() throws Exception {
	// Setup
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(130, 32);
	final int numExclusiveBuffers = 2;
	final int numFloatingBuffers = 128;

	final ExecutorService executor = Executors.newFixedThreadPool(2);

	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel  = createRemoteInputChannel(inputGate);
	inputGate.setInputChannel(inputChannel.partitionId.getPartitionId(), inputChannel);
	Throwable thrown = null;
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);
		inputChannel.requestSubpartition(0);

		final Callable<Void> requestBufferTask = new Callable<Void>() {
			@Override
			public Void call() throws Exception {
				while (true) {
					for (int j = 1; j <= numFloatingBuffers; j++) {
						inputChannel.onSenderBacklog(j);
					}

					if (inputChannel.isReleased()) {
						return null;
					}
				}
			}
		};

		final Callable<Void> releaseTask = new Callable<Void>() {
			@Override
			public Void call() throws Exception {
				inputChannel.releaseAllResources();

				return null;
			}
		};

		// Submit tasks and wait to finish
		submitTasksAndWaitForResults(executor, new Callable[]{requestBufferTask, releaseTask});

		assertEquals("There should be no buffers available in the channel.",
			0, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 130 buffers available in local pool.",
			130, bufferPool.getNumberOfAvailableMemorySegments() + networkBufferPool.getNumberOfAvailableMemorySegments());
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, executor, null, thrown, inputChannel);
	}
}
 

@Test
public void testIdleTime() throws IOException, InterruptedException {
	// setup
	final NetworkBufferPool globalPool = new NetworkBufferPool(10, 128, 2);
	final BufferPool localPool = globalPool.createBufferPool(1, 1, null, 1, Integer.MAX_VALUE);
	final ResultPartitionWriter resultPartition = new ResultPartitionBuilder()
		.setBufferPoolFactory(p -> localPool)
		.build();
	resultPartition.setup();
	final ResultPartitionWriter partitionWrapper = new ConsumableNotifyingResultPartitionWriterDecorator(
		new NoOpTaskActions(),
		new JobID(),
		resultPartition,
		new NoOpResultPartitionConsumableNotifier());
	final RecordWriter recordWriter = createRecordWriter(partitionWrapper);
	BufferBuilder builder = recordWriter.requestNewBufferBuilder(0);
	BufferBuilderTestUtils.fillBufferBuilder(builder, 1).finish();
	ResultSubpartitionView readView = resultPartition.getSubpartition(0).createReadView(new NoOpBufferAvailablityListener());
	Buffer buffer = readView.getNextBuffer().buffer();

	// idle time is zero when there is buffer available.
	assertEquals(0, recordWriter.getIdleTimeMsPerSecond().getCount());

	CountDownLatch syncLock = new CountDownLatch(1);
	AtomicReference<BufferBuilder> asyncRequestResult = new AtomicReference<>();
	final Thread requestThread = new Thread(new Runnable() {
		@Override
		public void run() {
			try {
				// notify that the request thread start to run.
				syncLock.countDown();
				// wait for buffer.
				asyncRequestResult.set(recordWriter.requestNewBufferBuilder(0));
			} catch (Exception e) {
			}
		}
	});
	requestThread.start();

	// wait until request thread start to run.
	syncLock.await();

	Thread.sleep(10);

	//recycle the buffer
	buffer.recycleBuffer();
	requestThread.join();

	assertThat(recordWriter.getIdleTimeMsPerSecond().getCount(), Matchers.greaterThan(0L));
	assertNotNull(asyncRequestResult.get());
}
 

/**
 * Tests a fix for FLINK-12544.
 *
 * @see <a href="https://issues.apache.org/jira/browse/FLINK-12544">FLINK-12544</a>
 */
@Test
public void testConcurrentRequestAndReleaseMemory() throws Exception {
	final ExecutorService executor = Executors.newFixedThreadPool(2);
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	try {
		final CountDownLatch blockLatch = new CountDownLatch(1);
		final CountDownLatch doneLatch = new CountDownLatch(1);

		final IOManager ioManager = new IOManagerAsyncWithCountDownLatch(blockLatch, doneLatch);
		final ResultPartitionWithCountDownLatch partition = new ResultPartitionWithCountDownLatch(
			"Test",
			new NoOpTaskActions(),
			new JobID(),
			new ResultPartitionID(),
			ResultPartitionType.BLOCKING,
			1,
			1,
			new ResultPartitionManager(),
			new NoOpResultPartitionConsumableNotifier(),
			ioManager,
			true,
			doneLatch,
			blockLatch);
		final BufferPool bufferPool = networkBufferPool.createBufferPool(1, 1, Optional.of(partition));
		partition.registerBufferPool(bufferPool);

		final BufferBuilder firstBuffer = bufferPool.requestBufferBuilderBlocking();
		partition.addBufferConsumer(firstBuffer.createBufferConsumer(), 0);
		// Finishes the buffer consumer which could be recycled during SpillableSubpartition#releaseMemory
		firstBuffer.finish();

		Future<Void> future = executor.submit(new Callable<Void>() {
			@Override
			public Void call() throws Exception {
				//Occupies the lock in SpillableSubpartition#releaseMemory, trying to get the lock in LocalBufferPool#recycle
				partition.releaseMemory(1);
				return null;
			}
		});

		final CompletableFuture<?> firstCallFuture = partition.getFirstCallFuture();
		firstCallFuture.thenRunAsync(() -> {
			try {
				// There are no available buffers in pool, so trigger release memory in SpillableSubpartition.
				// Occupies the lock in LocalBufferPool, and trying to get the lock in SpillableSubpartition.
				BufferBuilder secondBuffer = bufferPool.requestBufferBuilderBlocking();

				assertThat(firstBuffer, is(equalTo(secondBuffer)));
			} catch (IOException | InterruptedException ex) {
				fail("Should not throw any exceptions!");
			}
		}, executor);

		future.get();
	} finally {
		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
		executor.shutdown();
	}
}
 

/**
 * Tests to verify that there is no race condition with two things running in parallel:
 * requesting floating buffers on sender backlog and some other thread releasing
 * the input channel.
 */
@Test
public void testConcurrentOnSenderBacklogAndRelease() throws Exception {
	// Setup
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(130, 32, 2);
	final int numFloatingBuffers = 128;

	final ExecutorService executor = Executors.newFixedThreadPool(2);

	final SingleInputGate inputGate = createSingleInputGate(1, networkBufferPool);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	Throwable thrown = null;
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		final Callable<Void> requestBufferTask = new Callable<Void>() {
			@Override
			public Void call() throws Exception {
				while (true) {
					for (int j = 1; j <= numFloatingBuffers; j++) {
						inputChannel.onSenderBacklog(j);
					}

					if (inputChannel.isReleased()) {
						return null;
					}
				}
			}
		};

		final Callable<Void> releaseTask = new Callable<Void>() {
			@Override
			public Void call() throws Exception {
				inputChannel.releaseAllResources();

				return null;
			}
		};

		// Submit tasks and wait to finish
		submitTasksAndWaitForResults(executor, new Callable[]{requestBufferTask, releaseTask});

		assertEquals("There should be no buffers available in the channel.",
			0, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 130 buffers available in local pool.",
			130, bufferPool.getNumberOfAvailableMemorySegments() + networkBufferPool.getNumberOfAvailableMemorySegments());
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, executor, null, thrown, inputChannel);
	}
}
 

/**
 * Tests that failures are propagated correctly if
 * {@link RemoteInputChannel#notifyBufferAvailable(int)} throws an exception. Also tests that
 * a second listener will be notified in this case.
 */
@Test
public void testFailureInNotifyBufferAvailable() throws Exception {
	// Setup
	final int numExclusiveBuffers = 1;
	final int numFloatingBuffers = 1;
	final int numTotalBuffers = numExclusiveBuffers + numFloatingBuffers;
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(
		numTotalBuffers, 32, numExclusiveBuffers);

	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel successfulRemoteIC = createRemoteInputChannel(inputGate);
	successfulRemoteIC.requestSubpartition(0);

	// late creation -> no exclusive buffers, also no requested subpartition in successfulRemoteIC
	// (to trigger a failure in RemoteInputChannel#notifyBufferAvailable())
	final RemoteInputChannel failingRemoteIC = createRemoteInputChannel(inputGate);

	Buffer buffer = null;
	Throwable thrown = null;
	try {
		final BufferPool bufferPool =
			networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers);
		inputGate.setBufferPool(bufferPool);

		buffer = checkNotNull(bufferPool.requestBuffer());

		// trigger subscription to buffer pool
		failingRemoteIC.onSenderBacklog(1);
		successfulRemoteIC.onSenderBacklog(numExclusiveBuffers + 1);
		// recycling will call RemoteInputChannel#notifyBufferAvailable() which will fail and
		// this exception will be swallowed and set as an error in failingRemoteIC
		buffer.recycleBuffer();
		buffer = null;
		try {
			failingRemoteIC.checkError();
			fail("The input channel should have an error based on the failure in RemoteInputChannel#notifyBufferAvailable()");
		} catch (IOException e) {
			assertThat(e, hasProperty("cause", isA(IllegalStateException.class)));
		}
		// currently, the buffer is still enqueued in the bufferQueue of failingRemoteIC
		assertEquals(0, bufferPool.getNumberOfAvailableMemorySegments());
		buffer = successfulRemoteIC.requestBuffer();
		assertNull("buffer should still remain in failingRemoteIC", buffer);

		// releasing resources in failingRemoteIC should free the buffer again and immediately
		// recycle it into successfulRemoteIC
		failingRemoteIC.releaseAllResources();
		assertEquals(0, bufferPool.getNumberOfAvailableMemorySegments());
		buffer = successfulRemoteIC.requestBuffer();
		assertNotNull("no buffer given to successfulRemoteIC", buffer);
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, null, buffer, thrown, failingRemoteIC, successfulRemoteIC);
	}
}
 

/**
 * Verifies that {@link RemoteInputChannel} is enqueued in the pipeline for notifying credits,
 * and verifies the behaviour of credit notification by triggering channel's writability changed.
 */
@Test
public void testNotifyCreditAvailable() throws Exception {
	final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = new PartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32);
	final SingleInputGate inputGate = createSingleInputGate();
	final RemoteInputChannel inputChannel1 = createRemoteInputChannel(inputGate, client);
	final RemoteInputChannel inputChannel2 = createRemoteInputChannel(inputGate, client);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		final int numExclusiveBuffers = 2;
		inputGate.assignExclusiveSegments(networkBufferPool, numExclusiveBuffers);

		inputChannel1.requestSubpartition(0);
		inputChannel2.requestSubpartition(0);

		// The two input channels should send partition requests
		assertTrue(channel.isWritable());
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel1.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel2.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		// The buffer response will take one available buffer from input channel, and it will trigger
		// requesting (backlog + numExclusiveBuffers - numAvailableBuffers) floating buffers
		final BufferResponse bufferResponse1 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel1.getInputChannelId(), 1);
		final BufferResponse bufferResponse2 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel2.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse2);

		assertEquals(2, inputChannel1.getUnannouncedCredit());
		assertEquals(2, inputChannel2.getUnannouncedCredit());

		channel.runPendingTasks();

		// The two input channels should notify credits availability via the writable channel
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(inputChannel1.getInputChannelId(), ((AddCredit) readFromOutbound).receiverId);
		assertEquals(2, ((AddCredit) readFromOutbound).credit);

		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(inputChannel2.getInputChannelId(), ((AddCredit) readFromOutbound).receiverId);
		assertEquals(2, ((AddCredit) readFromOutbound).credit);
		assertNull(channel.readOutbound());

		ByteBuf channelBlockingBuffer = blockChannel(channel);

		// Trigger notify credits availability via buffer response on the condition of an un-writable channel
		final BufferResponse bufferResponse3 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 1, inputChannel1.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse3);

		assertEquals(1, inputChannel1.getUnannouncedCredit());
		assertEquals(0, inputChannel2.getUnannouncedCredit());

		channel.runPendingTasks();

		// The input channel will not notify credits via un-writable channel
		assertFalse(channel.isWritable());
		assertNull(channel.readOutbound());

		// Flush the buffer to make the channel writable again
		channel.flush();
		assertSame(channelBlockingBuffer, channel.readOutbound());

		// The input channel should notify credits via channel's writability changed event
		assertTrue(channel.isWritable());
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(1, ((AddCredit) readFromOutbound).credit);
		assertEquals(0, inputChannel1.getUnannouncedCredit());
		assertEquals(0, inputChannel2.getUnannouncedCredit());

		// no more messages
		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.releaseAllResources();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

/**
 * Verifies that {@link RemoteInputChannel} is enqueued in the pipeline for notifying credits,
 * and verifies the behaviour of credit notification by triggering channel's writability changed.
 */
@Test
public void testNotifyCreditAvailable() throws Exception {
	final CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = new NettyPartitionRequestClient(
		channel, handler, mock(ConnectionID.class), mock(PartitionRequestClientFactory.class));

	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, 32, 2);
	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel inputChannel1 = createRemoteInputChannel(inputGate, client, networkBufferPool);
	final RemoteInputChannel inputChannel2 = createRemoteInputChannel(inputGate, client, networkBufferPool);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();

		inputChannel1.requestSubpartition(0);
		inputChannel2.requestSubpartition(0);

		// The two input channels should send partition requests
		assertTrue(channel.isWritable());
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel1.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(inputChannel2.getInputChannelId(), ((PartitionRequest) readFromOutbound).receiverId);
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		// The buffer response will take one available buffer from input channel, and it will trigger
		// requesting (backlog + numExclusiveBuffers - numAvailableBuffers) floating buffers
		final BufferResponse bufferResponse1 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel1.getInputChannelId(), 1);
		final BufferResponse bufferResponse2 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel2.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse2);

		assertEquals(2, inputChannel1.getUnannouncedCredit());
		assertEquals(2, inputChannel2.getUnannouncedCredit());

		channel.runPendingTasks();

		// The two input channels should notify credits availability via the writable channel
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(inputChannel1.getInputChannelId(), ((AddCredit) readFromOutbound).receiverId);
		assertEquals(2, ((AddCredit) readFromOutbound).credit);

		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(inputChannel2.getInputChannelId(), ((AddCredit) readFromOutbound).receiverId);
		assertEquals(2, ((AddCredit) readFromOutbound).credit);
		assertNull(channel.readOutbound());

		ByteBuf channelBlockingBuffer = blockChannel(channel);

		// Trigger notify credits availability via buffer response on the condition of an un-writable channel
		final BufferResponse bufferResponse3 = createBufferResponse(
			TestBufferFactory.createBuffer(32), 1, inputChannel1.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse3);

		assertEquals(1, inputChannel1.getUnannouncedCredit());
		assertEquals(0, inputChannel2.getUnannouncedCredit());

		channel.runPendingTasks();

		// The input channel will not notify credits via un-writable channel
		assertFalse(channel.isWritable());
		assertNull(channel.readOutbound());

		// Flush the buffer to make the channel writable again
		channel.flush();
		assertSame(channelBlockingBuffer, channel.readOutbound());

		// The input channel should notify credits via channel's writability changed event
		assertTrue(channel.isWritable());
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(AddCredit.class));
		assertEquals(1, ((AddCredit) readFromOutbound).credit);
		assertEquals(0, inputChannel1.getUnannouncedCredit());
		assertEquals(0, inputChannel2.getUnannouncedCredit());

		// no more messages
		assertNull(channel.readOutbound());
	} finally {
		// Release all the buffer resources
		inputGate.close();

		networkBufferPool.destroyAllBufferPools();
		networkBufferPool.destroy();
	}
}
 

/**
 * Tests to verify that there is no race condition with two things running in parallel:
 * requesting floating buffers on sender backlog and some other thread releasing
 * the input channel.
 */
@Test
public void testConcurrentOnSenderBacklogAndRelease() throws Exception {
	// Setup
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(130, 32, 2);
	final int numFloatingBuffers = 128;

	final ExecutorService executor = Executors.newFixedThreadPool(2);

	final SingleInputGate inputGate = createSingleInputGate(1);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, networkBufferPool);
	Throwable thrown = null;
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		final Callable<Void> requestBufferTask = new Callable<Void>() {
			@Override
			public Void call() throws Exception {
				while (true) {
					for (int j = 1; j <= numFloatingBuffers; j++) {
						inputChannel.onSenderBacklog(j);
					}

					if (inputChannel.isReleased()) {
						return null;
					}
				}
			}
		};

		final Callable<Void> releaseTask = new Callable<Void>() {
			@Override
			public Void call() throws Exception {
				inputChannel.releaseAllResources();

				return null;
			}
		};

		// Submit tasks and wait to finish
		submitTasksAndWaitForResults(executor, new Callable[]{requestBufferTask, releaseTask});

		assertEquals("There should be no buffers available in the channel.",
			0, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be 130 buffers available in local pool.",
			130, bufferPool.getNumberOfAvailableMemorySegments() + networkBufferPool.getNumberOfAvailableMemorySegments());
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, executor, null, thrown, inputChannel);
	}
}
 

/**
 * Verifies that {@link BufferResponse} of compressed {@link Buffer} can be handled correctly.
 */
@Test
public void testReceiveCompressedBuffer() throws Exception {
	int bufferSize = 1024;
	String compressionCodec = "LZ4";
	BufferCompressor compressor = new BufferCompressor(bufferSize, compressionCodec);
	BufferDecompressor decompressor = new BufferDecompressor(bufferSize, compressionCodec);
	NetworkBufferPool networkBufferPool = new NetworkBufferPool(10, bufferSize, 2);
	SingleInputGate inputGate = new SingleInputGateBuilder()
		.setBufferDecompressor(decompressor)
		.setSegmentProvider(networkBufferPool)
		.build();
	RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate, null);
	inputGate.setInputChannels(inputChannel);

	try {
		BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();

		CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
		handler.addInputChannel(inputChannel);

		Buffer buffer = compressor.compressToOriginalBuffer(TestBufferFactory.createBuffer(bufferSize));
		BufferResponse bufferResponse = createBufferResponse(
			buffer,
			0,
			inputChannel.getInputChannelId(),
			2,
			new NetworkBufferAllocator(handler));
		assertTrue(bufferResponse.isCompressed);
		handler.channelRead(null, bufferResponse);

		Buffer receivedBuffer = inputChannel.getNextReceivedBuffer();
		assertNotNull(receivedBuffer);
		assertTrue(receivedBuffer.isCompressed());
		receivedBuffer.recycleBuffer();
	} finally {
		releaseResource(inputGate, networkBufferPool);
	}
}
 

/**
 * Tests to verify that there is no race condition with two things running in parallel:
 * recycling the exclusive or floating buffers and some other thread releasing the
 * input channel.
 */
@Test
public void testConcurrentRecycleAndRelease() throws Exception {
	// Setup
	final int numExclusiveSegments = 120;
	final NetworkBufferPool networkBufferPool = new NetworkBufferPool(248, 32, numExclusiveSegments);
	final int numFloatingBuffers = 128;

	final ExecutorService executor = Executors.newFixedThreadPool(3);

	final SingleInputGate inputGate = createSingleInputGate(1, networkBufferPool);
	final RemoteInputChannel inputChannel = createRemoteInputChannel(inputGate);
	inputGate.setInputChannels(inputChannel);
	Throwable thrown = null;
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(numFloatingBuffers, numFloatingBuffers);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();
		inputChannel.requestSubpartition(0);

		final Callable<Void> releaseTask = new Callable<Void>() {
			@Override
			public Void call() throws Exception {
				inputChannel.releaseAllResources();

				return null;
			}
		};

		// Submit tasks and wait to finish
		submitTasksAndWaitForResults(executor, new Callable[]{
			recycleBufferTask(inputChannel, bufferPool, numExclusiveSegments, numFloatingBuffers),
			releaseTask});

		assertEquals("There should be no buffers available in the channel.",
			0, inputChannel.getNumberOfAvailableBuffers());
		assertEquals("There should be " + numFloatingBuffers + " buffers available in local pool.",
			numFloatingBuffers, bufferPool.getNumberOfAvailableMemorySegments());
		assertEquals("There should be " + numExclusiveSegments + " buffers available in global pool.",
			numExclusiveSegments, networkBufferPool.getNumberOfAvailableMemorySegments());
	} catch (Throwable t) {
		thrown = t;
	} finally {
		cleanup(networkBufferPool, executor, null, thrown, inputChannel);
	}
}