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

/**
 * Tests that the buffer is recycled correctly if exception is thrown during
 * {@link ChannelStateReader#readOutputData(ResultSubpartitionInfo, BufferBuilder)}.
 */
@Test
public void testReadRecoveredStateWithException() throws Exception {
	final int totalBuffers = 2;
	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, 1, 1);
	final ResultPartition partition = new ResultPartitionBuilder()
		.setNetworkBufferPool(globalPool)
		.build();
	final ChannelStateReader stateReader = new ChannelStateReaderWithException();

	try {
		partition.setup();
		partition.readRecoveredState(stateReader);
	} catch (IOException e) {
		assertThat("should throw custom exception message", e.getMessage().contains("test"));
	} finally {
		globalPool.destroyAllBufferPools();
		// verify whether there are any buffers leak
		assertEquals(totalBuffers, globalPool.getNumberOfAvailableMemorySegments());
		globalPool.destroy();
	}
}
 

@Test
public void testInitializeEmptyState() throws Exception {
	final int totalBuffers = 2;
	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, 1, 1);
	final ResultPartition partition = new ResultPartitionBuilder()
		.setNetworkBufferPool(globalPool)
		.build();
	final ChannelStateReader stateReader = ChannelStateReader.NO_OP;
	try {
		partition.setup();
		partition.readRecoveredState(stateReader);

		for (ResultSubpartition subpartition : partition.getAllPartitions()) {
			// no buffers are added into the queue for empty states
			assertEquals(0, subpartition.getTotalNumberOfBuffers());
		}

		// destroy the local pool to verify that all the requested buffers by partition are recycled
		partition.getBufferPool().lazyDestroy();
		assertEquals(totalBuffers, globalPool.getNumberOfAvailableMemorySegments());
	} finally {
		// cleanup
		globalPool.destroyAllBufferPools();
		globalPool.destroy();
	}
}
 

private void testReadEmptyStateOrThrowException(boolean isRemote, ChannelStateReader reader) throws Exception {
	// setup
	final int totalBuffers = 10;
	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, 32, 2);
	final SingleInputGate inputGate = createInputGate(globalPool);
	final RecoveredInputChannel inputChannel = createRecoveredChannel(isRemote, inputGate);

	try {
		inputGate.setInputChannels(inputChannel);
		inputGate.setup();

		// it would throw expected exception for the case of testReadStateWithException.
		inputChannel.readRecoveredState(reader);

		// the channel only has one EndOfChannelStateEvent in the queue for the case of testReadEmptyState.
		assertEquals(1, inputChannel.getNumberOfQueuedBuffers());
		assertFalse(inputChannel.getNextBuffer().isPresent());
		assertTrue(inputChannel.getStateConsumedFuture().isDone());
	} finally {
		// cleanup and verify no buffer leak
		inputGate.close();
		globalPool.destroyAllBufferPools();
		assertEquals(totalBuffers, globalPool.getNumberOfAvailableMemorySegments());
		globalPool.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();
	}
}
 

private void testPartitionBufferPool(ResultPartitionType type) throws Exception {
	//setup
	final int networkBuffersPerChannel = 2;
	final int floatingNetworkBuffersPerGate = 8;
	final NetworkBufferPool globalPool = new NetworkBufferPool(20, 1, 1);
	final ResultPartition partition = new ResultPartitionBuilder()
		.setResultPartitionType(type)
		.setFileChannelManager(fileChannelManager)
		.setNetworkBuffersPerChannel(networkBuffersPerChannel)
		.setFloatingNetworkBuffersPerGate(floatingNetworkBuffersPerGate)
		.setNetworkBufferPool(globalPool)
		.build();

	try {
		partition.setup();
		BufferPool bufferPool = partition.getBufferPool();
		// verify the amount of buffers in created local pool
		assertEquals(partition.getNumberOfSubpartitions() + 1, bufferPool.getNumberOfRequiredMemorySegments());
		if (type.isBounded()) {
			final int maxNumBuffers = networkBuffersPerChannel * partition.getNumberOfSubpartitions() + floatingNetworkBuffersPerGate;
			assertEquals(maxNumBuffers, bufferPool.getMaxNumberOfMemorySegments());
		} else {
			assertEquals(Integer.MAX_VALUE, bufferPool.getMaxNumberOfMemorySegments());
		}

	} finally {
		// cleanup
		globalPool.destroyAllBufferPools();
		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 CreditBasedPartitionRequestClientHandler handler = new CreditBasedPartitionRequestClientHandler();
	final EmbeddedChannel channel = new EmbeddedChannel(handler);
	final PartitionRequestClient client = createPartitionRequestClient(channel, handler);

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

	try {
		inputGate.setInputChannels(inputChannel);
		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();
	}
}
 

/**
 * 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();
	}
}
 

/**
 * 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);
	final RemoteInputChannel inputChannel = InputChannelBuilder.newBuilder()
		.setMemorySegmentProvider(networkBufferPool)
		.buildRemoteAndSetToGate(inputGate);
	try {
		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);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

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

		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);
	final RemoteInputChannel inputChannel = InputChannelBuilder.newBuilder()
		.setMemorySegmentProvider(networkBufferPool)
		.buildRemoteAndSetToGate(inputGate);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(8, 8);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();

		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.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();
	}
}
 

/**
 * 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} 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();
	}
}
 

/**
 * Verifies that {@link RemoteInputChannel} is enqueued in the pipeline, but {@link AddCredit}
 * message is not sent actually when this input channel is released.
 */
@Test
public void testNotifyCreditAvailableAfterReleased() 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 inputChannel = createRemoteInputChannel(inputGate, client, networkBufferPool);
	try {
		final BufferPool bufferPool = networkBufferPool.createBufferPool(6, 6);
		inputGate.setBufferPool(bufferPool);
		inputGate.assignExclusiveSegments();

		inputChannel.requestSubpartition(0);

		// This should send the partition request
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		// Trigger request floating buffers via buffer response to notify credits available
		final BufferResponse bufferResponse = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(2, inputChannel.getUnannouncedCredit());

		// Release the input channel
		inputGate.close();

		// it should send a close request after releasing the input channel,
		// but will not notify credits for a released input channel.
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(CloseRequest.class));

		channel.runPendingTasks();

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

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

@Test
public void testEmitRecordWithPartitionStateRecovery() throws Exception {
	final int totalBuffers = 10; // enough for both states and normal records
	final int totalStates = 2;
	final int[] states = {1, 2, 3, 4};
	final int[] records = {5, 6, 7, 8};
	final int bufferSize = states.length * Integer.BYTES;

	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, bufferSize, 1);
	final ChannelStateReader stateReader = new ResultPartitionTest.FiniteChannelStateReader(totalStates, states);
	final ResultPartition partition = new ResultPartitionBuilder()
		.setNetworkBufferPool(globalPool)
		.build();
	final RecordWriter<IntValue> recordWriter = new RecordWriterBuilder<IntValue>().build(partition);

	try {
		partition.setup();
		partition.readRecoveredState(stateReader);

		for (int record: records) {
			// the record length 4 is also written into buffer for every emit
			recordWriter.broadcastEmit(new IntValue(record));
		}

		// every buffer can contain 2 int records with 2 int length(4)
		final int[][] expectedRecordsInBuffer = {{4, 5, 4, 6}, {4, 7, 4, 8}};

		for (ResultSubpartition subpartition : partition.getAllPartitions()) {
			// create the view to consume all the buffers with states and records
			final ResultSubpartitionView view = new PipelinedSubpartitionView(
				(PipelinedSubpartition) subpartition,
				new NoOpBufferAvailablityListener());

			int numConsumedBuffers = 0;
			ResultSubpartition.BufferAndBacklog bufferAndBacklog;
			while ((bufferAndBacklog = view.getNextBuffer()) != null) {
				Buffer buffer = bufferAndBacklog.buffer();
				int[] expected = numConsumedBuffers < totalStates ? states : expectedRecordsInBuffer[numConsumedBuffers - totalStates];
				BufferBuilderAndConsumerTest.assertContent(
					buffer,
					partition.getBufferPool()
						.getSubpartitionBufferRecyclers()[subpartition.getSubPartitionIndex()],
					expected);

				buffer.recycleBuffer();
				numConsumedBuffers++;
			}

			assertEquals(totalStates + expectedRecordsInBuffer.length, numConsumedBuffers);
		}
	} finally {
		// cleanup
		globalPool.destroyAllBufferPools();
		globalPool.destroy();
	}
}
 

/**
 * 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();
	}
}
 

@Test
public void testInitializeMoreStateThanBuffer() throws Exception {
	final int totalBuffers = 2; // the total buffers are less than the requirement from total states
	final int totalStates = 5;
	final int[] states = {1, 2, 3, 4};
	final int bufferSize = states.length * Integer.BYTES;

	final NetworkBufferPool globalPool = new NetworkBufferPool(totalBuffers, bufferSize, 1);
	final ChannelStateReader stateReader = new FiniteChannelStateReader(totalStates, states);
	final ResultPartition partition = new ResultPartitionBuilder()
		.setNetworkBufferPool(globalPool)
		.build();
	final ExecutorService executor = Executors.newFixedThreadPool(1);

	try {
		final Callable<Void> partitionConsumeTask = () -> {
			for (ResultSubpartition subpartition : partition.getAllPartitions()) {
				final ResultSubpartitionView view = new PipelinedSubpartitionView(
					(PipelinedSubpartition) subpartition,
					new NoOpBufferAvailablityListener());

				int numConsumedBuffers = 0;
				while (numConsumedBuffers != totalStates) {
					ResultSubpartition.BufferAndBacklog bufferAndBacklog = view.getNextBuffer();
					if (bufferAndBacklog != null) {
						Buffer buffer = bufferAndBacklog.buffer();
						BufferBuilderAndConsumerTest.assertContent(
							buffer,
							partition.getBufferPool()
								.getSubpartitionBufferRecyclers()[subpartition.getSubPartitionIndex()],
							states);
						buffer.recycleBuffer();
						numConsumedBuffers++;
					} else {
						Thread.sleep(5);
					}
				}
				assertNull(view.getNextBuffer());
			}
			return null;
		};
		Future<Void> result = executor.submit(partitionConsumeTask);

		partition.setup();
		partition.readRecoveredState(stateReader);

		// wait the partition consume task finish
		result.get(20, TimeUnit.SECONDS);

		// destroy the local pool to verify that all the requested buffers by partition are recycled
		partition.getBufferPool().lazyDestroy();
		assertEquals(totalBuffers, globalPool.getNumberOfAvailableMemorySegments());
	} finally {
		// cleanup
		executor.shutdown();
		globalPool.destroyAllBufferPools();
		globalPool.destroy();
	}
}
 

/**
 * Tests the consumption of multiple subpartitions via local input channels.
 *
 * <p>Multiple producer tasks produce pipelined partitions, which are consumed by multiple
 * tasks via local input channels.
 */
@Test
public void testConcurrentConsumeMultiplePartitions() throws Exception {
	// Config
	final int parallelism = 32;
	final int producerBufferPoolSize = parallelism + 1;
	final int numberOfBuffersPerChannel = 1024;

	checkArgument(parallelism >= 1);
	checkArgument(producerBufferPoolSize >= parallelism);
	checkArgument(numberOfBuffersPerChannel >= 1);

	// Setup
	// One thread per produced partition and one per consumer
	final ExecutorService executor = Executors.newFixedThreadPool(2 * parallelism);

	final NetworkBufferPool networkBuffers = new NetworkBufferPool(
		(parallelism * producerBufferPoolSize) + (parallelism * parallelism),
		TestBufferFactory.BUFFER_SIZE, 1);

	final ResultPartitionManager partitionManager = new ResultPartitionManager();

	final ResultPartitionID[] partitionIds = new ResultPartitionID[parallelism];
	final TestPartitionProducer[] partitionProducers = new TestPartitionProducer[parallelism];

	// Create all partitions
	for (int i = 0; i < parallelism; i++) {
		partitionIds[i] = new ResultPartitionID();

		final ResultPartition partition = new ResultPartitionBuilder()
			.setResultPartitionId(partitionIds[i])
			.setNumberOfSubpartitions(parallelism)
			.setNumTargetKeyGroups(parallelism)
			.setResultPartitionManager(partitionManager)
			.setBufferPoolFactory(p ->
				networkBuffers.createBufferPool(producerBufferPoolSize, producerBufferPoolSize))
			.build();

		// Create a buffer pool for this partition
		partition.setup();

		// Create the producer
		partitionProducers[i] = new TestPartitionProducer(
			partition,
			false,
			new TestPartitionProducerBufferSource(
				parallelism,
				partition.getBufferPool(),
				numberOfBuffersPerChannel)
		);
	}

	// Test
	try {
		// Submit producer tasks
		List<CompletableFuture<?>> results = Lists.newArrayListWithCapacity(
			parallelism + 1);

		for (int i = 0; i < parallelism; i++) {
			results.add(CompletableFuture.supplyAsync(
				CheckedSupplier.unchecked(partitionProducers[i]::call), executor));
		}

		// Submit consumer
		for (int i = 0; i < parallelism; i++) {
			final TestLocalInputChannelConsumer consumer = new TestLocalInputChannelConsumer(
				i,
				parallelism,
				numberOfBuffersPerChannel,
				networkBuffers.createBufferPool(parallelism, parallelism),
				partitionManager,
				new TaskEventDispatcher(),
				partitionIds);

			results.add(CompletableFuture.supplyAsync(CheckedSupplier.unchecked(consumer::call), executor));
		}

		FutureUtils.waitForAll(results).get();
	}
	finally {
		networkBuffers.destroyAllBufferPools();
		networkBuffers.destroy();
		executor.shutdown();
	}
}
 

/**
 * Verifies that {@link RemoteInputChannel} is enqueued in the pipeline, but {@link AddCredit}
 * message is not sent actually when this input channel is released.
 */
@Test
public void testNotifyCreditAvailableAfterReleased() 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 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);

		// This should send the partition request
		Object readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(PartitionRequest.class));
		assertEquals(2, ((PartitionRequest) readFromOutbound).credit);

		// Trigger request floating buffers via buffer response to notify credits available
		final BufferResponse bufferResponse = createBufferResponse(
			TestBufferFactory.createBuffer(32), 0, inputChannel.getInputChannelId(), 1);
		handler.channelRead(mock(ChannelHandlerContext.class), bufferResponse);

		assertEquals(2, inputChannel.getUnannouncedCredit());

		// Release the input channel
		inputGate.releaseAllResources();

		// it should send a close request after releasing the input channel,
		// but will not notify credits for a released input channel.
		readFromOutbound = channel.readOutbound();
		assertThat(readFromOutbound, instanceOf(CloseRequest.class));

		channel.runPendingTasks();

		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 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();
	}
}