You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@flink.apache.org by gu...@apache.org on 2023/04/04 02:35:29 UTC

[flink] branch master updated (7c6d8b0134c -> fb6caee1371)

This is an automated email from the ASF dual-hosted git repository.

guoweijie pushed a change to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git


    from 7c6d8b0134c [FLINK-31707][python] Fix Pandas UDAF to support accepting constant string as inputs
     new 0542fae7e41 [hotfix][test] Migrate LocalBufferPoolTest to Junit5 and AssertJ.
     new 0a49f9d1e88 [hotfix] Some minor clean-up and typo fix for LocalBufferPool.
     new cfe393dd870 [hotfix] Add missing @GuardedBy annotation for LocalBufferPool.
     new fb6caee1371 [FLINK-31293][runtime] LocalBufferPool request overdraft buffer only when no available buffer and pool size is reached.

The 4 revisions listed above as "new" are entirely new to this
repository and will be described in separate emails.  The revisions
listed as "add" were already present in the repository and have only
been added to this reference.


Summary of changes:
 .../runtime/io/network/buffer/LocalBufferPool.java |  60 ++-
 .../io/network/buffer/LocalBufferPoolTest.java     | 480 ++++++++++++++-------
 .../io/network/buffer/TestingBufferListener.java   |  73 ++++
 3 files changed, 442 insertions(+), 171 deletions(-)
 create mode 100644 flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/TestingBufferListener.java


[flink] 01/04: [hotfix][test] Migrate LocalBufferPoolTest to Junit5 and AssertJ.

Posted by gu...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

guoweijie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit 0542fae7e412ef1447cd3a174013a3743bdaaa30
Author: Weijie Guo <re...@163.com>
AuthorDate: Tue Mar 7 01:05:47 2023 +0800

    [hotfix][test] Migrate LocalBufferPoolTest to Junit5 and AssertJ.
---
 .../io/network/buffer/LocalBufferPoolTest.java     | 302 +++++++++++----------
 .../io/network/buffer/TestingBufferListener.java   |  73 +++++
 2 files changed, 227 insertions(+), 148 deletions(-)

diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
index 3149697471d..c677cd4905f 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
@@ -22,15 +22,13 @@ import org.apache.flink.core.fs.AutoCloseableRegistry;
 import org.apache.flink.core.memory.MemorySegment;
 import org.apache.flink.core.testutils.CheckedThread;
 import org.apache.flink.runtime.execution.CancelTaskException;
-import org.apache.flink.testutils.executor.TestExecutorResource;
-import org.apache.flink.util.TestLogger;
+import org.apache.flink.testutils.executor.TestExecutorExtension;
 
-import org.junit.After;
-import org.junit.Before;
-import org.junit.ClassRule;
-import org.junit.Test;
+import org.junit.jupiter.api.AfterEach;
+import org.junit.jupiter.api.BeforeEach;
+import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.Timeout;
-import org.mockito.Mockito;
+import org.junit.jupiter.api.extension.RegisterExtension;
 
 import javax.annotation.Nullable;
 
@@ -52,17 +50,11 @@ import java.util.concurrent.atomic.AtomicInteger;
 
 import static org.apache.flink.util.Preconditions.checkArgument;
 import static org.apache.flink.util.Preconditions.checkNotNull;
-import static org.junit.Assert.assertEquals;
-import static org.junit.Assert.assertFalse;
-import static org.junit.Assert.assertNotNull;
-import static org.junit.Assert.assertNull;
-import static org.junit.Assert.assertThrows;
-import static org.junit.Assert.assertTrue;
-import static org.mockito.Mockito.times;
-import static org.mockito.Mockito.verify;
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
 /** Tests for the {@link LocalBufferPool}. */
-public class LocalBufferPoolTest extends TestLogger {
+class LocalBufferPoolTest {
 
     private static final int numBuffers = 1024;
 
@@ -72,67 +64,73 @@ public class LocalBufferPoolTest extends TestLogger {
 
     private BufferPool localBufferPool;
 
-    @ClassRule
-    public static final TestExecutorResource<ExecutorService> EXECUTOR_RESOURCE =
-            new TestExecutorResource<>(() -> Executors.newCachedThreadPool());
+    @RegisterExtension
+    public static final TestExecutorExtension<ExecutorService> EXECUTOR_RESOURCE =
+            new TestExecutorExtension<>(Executors::newCachedThreadPool);
 
-    @Before
-    public void setupLocalBufferPool() throws Exception {
+    @BeforeEach
+    void setupLocalBufferPool() {
         networkBufferPool = new NetworkBufferPool(numBuffers, memorySegmentSize);
         localBufferPool = new LocalBufferPool(networkBufferPool, 1);
 
-        assertEquals(1, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isEqualTo(1);
     }
 
-    @After
-    public void destroyAndVerifyAllBuffersReturned() {
+    @AfterEach
+    void destroyAndVerifyAllBuffersReturned() {
         if (!localBufferPool.isDestroyed()) {
             localBufferPool.lazyDestroy();
         }
 
         String msg = "Did not return all buffers to memory segment pool after test.";
-        assertEquals(msg, numBuffers, networkBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(networkBufferPool.getNumberOfAvailableMemorySegments())
+                .withFailMessage(msg)
+                .isEqualTo(numBuffers);
         // no other local buffer pools used than the one above, but call just in case
         networkBufferPool.destroyAllBufferPools();
         networkBufferPool.destroy();
     }
 
     @Test
-    public void testReserveSegments() throws Exception {
+    void testReserveSegments() throws Exception {
         NetworkBufferPool networkBufferPool =
                 new NetworkBufferPool(2, memorySegmentSize, Duration.ofSeconds(2));
         try {
             BufferPool bufferPool1 = networkBufferPool.createBufferPool(1, 2);
-            assertThrows(IllegalArgumentException.class, () -> bufferPool1.reserveSegments(2));
+            assertThatThrownBy(() -> bufferPool1.reserveSegments(2))
+                    .isInstanceOf(IllegalArgumentException.class);
 
             // request all buffers
             ArrayList<Buffer> buffers = new ArrayList<>(2);
             buffers.add(bufferPool1.requestBuffer());
             buffers.add(bufferPool1.requestBuffer());
-            assertEquals(2, buffers.size());
+            assertThat(buffers).hasSize(2);
 
             BufferPool bufferPool2 = networkBufferPool.createBufferPool(1, 10);
-            assertThrows(IOException.class, () -> bufferPool2.reserveSegments(1));
-            assertFalse(bufferPool2.isAvailable());
+            assertThatThrownBy(() -> bufferPool2.reserveSegments(1))
+                    .isInstanceOf(IOException.class);
+            assertThat(bufferPool2.isAvailable()).isFalse();
 
             buffers.forEach(Buffer::recycleBuffer);
             bufferPool1.lazyDestroy();
             bufferPool2.lazyDestroy();
 
             BufferPool bufferPool3 = networkBufferPool.createBufferPool(2, 10);
-            assertEquals(1, bufferPool3.getNumberOfAvailableMemorySegments());
+            assertThat(bufferPool3.getNumberOfAvailableMemorySegments()).isEqualTo(1);
             bufferPool3.reserveSegments(2);
-            assertEquals(2, bufferPool3.getNumberOfAvailableMemorySegments());
+            assertThat(bufferPool3.getNumberOfAvailableMemorySegments()).isEqualTo(2);
 
             bufferPool3.lazyDestroy();
-            assertThrows(CancelTaskException.class, () -> bufferPool3.reserveSegments(1));
+            assertThatThrownBy(() -> bufferPool3.reserveSegments(1))
+                    .isInstanceOf(CancelTaskException.class);
         } finally {
             networkBufferPool.destroy();
         }
     }
 
-    @Test(timeout = 10000) // timeout can indicate a potential deadlock
-    public void testReserveSegmentsAndCancel() throws Exception {
+    @Test
+    @Timeout(10) // timeout can indicate a potential deadlock
+    void testReserveSegmentsAndCancel() throws Exception {
         int totalSegments = 4;
         int segmentsToReserve = 2;
 
@@ -185,22 +183,23 @@ public class LocalBufferPoolTest extends TestLogger {
         } finally {
             segments.forEach(localPool1::recycle);
             localPool1.lazyDestroy();
-            assertEquals(0, globalPool.getNumberOfUsedMemorySegments());
+            assertThat(globalPool.getNumberOfUsedMemorySegments()).isZero();
             globalPool.destroy();
         }
     }
 
     @Test
-    public void testRequestMoreThanAvailable() {
+    void testRequestMoreThanAvailable() {
         localBufferPool.setNumBuffers(numBuffers);
 
-        List<Buffer> requests = new ArrayList<Buffer>(numBuffers);
+        List<Buffer> requests = new ArrayList<>(numBuffers);
 
         for (int i = 1; i <= numBuffers; i++) {
             Buffer buffer = localBufferPool.requestBuffer();
 
-            assertEquals(Math.min(i + 1, numBuffers), getNumRequestedFromMemorySegmentPool());
-            assertNotNull(buffer);
+            assertThat(getNumRequestedFromMemorySegmentPool())
+                    .isEqualTo(Math.min(i + 1, numBuffers));
+            assertThat(buffer).isNotNull();
 
             requests.add(buffer);
         }
@@ -208,8 +207,8 @@ public class LocalBufferPoolTest extends TestLogger {
         {
             // One more...
             Buffer buffer = localBufferPool.requestBuffer();
-            assertEquals(numBuffers, getNumRequestedFromMemorySegmentPool());
-            assertNull(buffer);
+            assertThat(getNumRequestedFromMemorySegmentPool()).isEqualTo(numBuffers);
+            assertThat(buffer).isNull();
         }
 
         for (Buffer buffer : requests) {
@@ -218,22 +217,22 @@ public class LocalBufferPoolTest extends TestLogger {
     }
 
     @Test
-    public void testSetNumAfterDestroyDoesNotProactivelyFetchSegments() {
+    void testSetNumAfterDestroyDoesNotProactivelyFetchSegments() {
         localBufferPool.setNumBuffers(2);
-        assertEquals(2L, localBufferPool.getNumBuffers());
-        assertEquals(1L, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumBuffers()).isEqualTo(2);
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isOne();
 
         localBufferPool.lazyDestroy();
         localBufferPool.setNumBuffers(3);
-        assertEquals(3L, localBufferPool.getNumBuffers());
-        assertEquals(0L, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumBuffers()).isEqualTo(3);
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isZero();
     }
 
     @Test
-    public void testRecycleAfterDestroy() {
+    void testRecycleAfterDestroy() {
         localBufferPool.setNumBuffers(numBuffers);
 
-        List<Buffer> requests = new ArrayList<Buffer>(numBuffers);
+        List<Buffer> requests = new ArrayList<>(numBuffers);
 
         for (int i = 0; i < numBuffers; i++) {
             requests.add(localBufferPool.requestBuffer());
@@ -242,7 +241,7 @@ public class LocalBufferPoolTest extends TestLogger {
         localBufferPool.lazyDestroy();
 
         // All buffers have been requested, but can not be returned yet.
-        assertEquals(numBuffers, getNumRequestedFromMemorySegmentPool());
+        assertThat(getNumRequestedFromMemorySegmentPool()).isEqualTo(numBuffers);
 
         // Recycle should return buffers to memory segment pool
         for (Buffer buffer : requests) {
@@ -252,7 +251,7 @@ public class LocalBufferPoolTest extends TestLogger {
 
     @Test
     @Timeout(30)
-    public void testRequestBuffersOnRecycle() throws Exception {
+    void testRequestBuffersOnRecycle() throws Exception {
         BufferPool bufferPool1 = networkBufferPool.createBufferPool(512, 2048);
         List<MemorySegment> segments = new ArrayList<>();
         for (int i = 0; i < 1023; i++) {
@@ -282,27 +281,27 @@ public class LocalBufferPoolTest extends TestLogger {
     }
 
     @Test
-    public void testRecycleExcessBuffersAfterRecycling() {
+    void testRecycleExcessBuffersAfterRecycling() {
         localBufferPool.setNumBuffers(numBuffers);
 
-        List<Buffer> requests = new ArrayList<Buffer>(numBuffers);
+        List<Buffer> requests = new ArrayList<>(numBuffers);
 
         // Request all buffers
         for (int i = 1; i <= numBuffers; i++) {
             requests.add(localBufferPool.requestBuffer());
         }
 
-        assertEquals(numBuffers, getNumRequestedFromMemorySegmentPool());
+        assertThat(getNumRequestedFromMemorySegmentPool()).isEqualTo(numBuffers);
 
         // Reduce the number of buffers in the local pool
         localBufferPool.setNumBuffers(numBuffers / 2);
 
         // Need to wait until we recycle the buffers
-        assertEquals(numBuffers, getNumRequestedFromMemorySegmentPool());
+        assertThat(getNumRequestedFromMemorySegmentPool()).isEqualTo(numBuffers);
 
         for (int i = 1; i < numBuffers / 2; i++) {
             requests.remove(0).recycleBuffer();
-            assertEquals(numBuffers - i, getNumRequestedFromMemorySegmentPool());
+            assertThat(getNumRequestedFromMemorySegmentPool()).isEqualTo(numBuffers - i);
         }
 
         for (Buffer buffer : requests) {
@@ -311,10 +310,10 @@ public class LocalBufferPoolTest extends TestLogger {
     }
 
     @Test
-    public void testRecycleExcessBuffersAfterChangingNumBuffers() {
+    void testRecycleExcessBuffersAfterChangingNumBuffers() {
         localBufferPool.setNumBuffers(numBuffers);
 
-        List<Buffer> requests = new ArrayList<Buffer>(numBuffers);
+        List<Buffer> requests = new ArrayList<>(numBuffers);
 
         // Request all buffers
         for (int i = 1; i <= numBuffers; i++) {
@@ -326,18 +325,19 @@ public class LocalBufferPoolTest extends TestLogger {
             buffer.recycleBuffer();
         }
 
-        assertEquals(numBuffers, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isEqualTo(numBuffers);
 
         localBufferPool.setNumBuffers(numBuffers / 2);
 
-        assertEquals(numBuffers / 2, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isEqualTo(numBuffers / 2);
     }
 
-    @Test(expected = IllegalArgumentException.class)
-    public void testSetLessThanRequiredNumBuffers() {
+    @Test
+    void testSetLessThanRequiredNumBuffers() {
         localBufferPool.setNumBuffers(1);
 
-        localBufferPool.setNumBuffers(0);
+        assertThatThrownBy(() -> localBufferPool.setNumBuffers(0))
+                .isInstanceOf(IllegalArgumentException.class);
     }
 
     // ------------------------------------------------------------------------
@@ -345,38 +345,44 @@ public class LocalBufferPoolTest extends TestLogger {
     // ------------------------------------------------------------------------
 
     @Test
-    public void testPendingRequestWithListenersAfterRecycle() {
+    void testPendingRequestWithListenersAfterRecycle() {
         CountBufferListener listener1 = new CountBufferListener();
         CountBufferListener listener2 = new CountBufferListener();
 
         Buffer available = localBufferPool.requestBuffer();
 
-        assertNull(localBufferPool.requestBuffer());
+        assertThat(localBufferPool.requestBuffer()).isNull();
 
-        assertTrue(localBufferPool.addBufferListener(listener1));
-        assertTrue(localBufferPool.addBufferListener(listener2));
+        assertThat(localBufferPool.addBufferListener(listener1)).isTrue();
+        assertThat(localBufferPool.addBufferListener(listener2)).isTrue();
 
         // Recycle the buffer to notify both of the above listeners once
         checkNotNull(available).recycleBuffer();
 
-        assertEquals(1, listener1.getCount());
-        assertEquals(1, listener1.getCount());
+        assertThat(listener1.getCount()).isOne();
+        assertThat(listener1.getCount()).isOne();
 
-        assertFalse(localBufferPool.addBufferListener(listener1));
-        assertFalse(localBufferPool.addBufferListener(listener2));
+        assertThat(localBufferPool.addBufferListener(listener1)).isFalse();
+        assertThat(localBufferPool.addBufferListener(listener2)).isFalse();
     }
 
     @Test
-    @SuppressWarnings("unchecked")
-    public void testCancelPendingRequestsAfterDestroy() {
-        BufferListener listener = Mockito.mock(BufferListener.class);
+    void testCancelPendingRequestsAfterDestroy() {
+        AtomicInteger invokeNotifyBufferDestroyedCounter = new AtomicInteger(0);
+
+        BufferListener listener =
+                TestingBufferListener.builder()
+                        .setNotifyBufferDestroyedRunnable(
+                                invokeNotifyBufferDestroyedCounter::incrementAndGet)
+                        .build();
 
         localBufferPool.setNumBuffers(1);
 
         Buffer available = localBufferPool.requestBuffer();
         Buffer unavailable = localBufferPool.requestBuffer();
 
-        assertNull(unavailable);
+        assertThat(available).isNotNull();
+        assertThat(unavailable).isNull();
 
         localBufferPool.addBufferListener(listener);
 
@@ -384,7 +390,7 @@ public class LocalBufferPoolTest extends TestLogger {
 
         available.recycleBuffer();
 
-        verify(listener, times(1)).notifyBufferDestroyed();
+        assertThat(invokeNotifyBufferDestroyedCounter).hasValue(1);
     }
 
     // ------------------------------------------------------------------------
@@ -393,7 +399,7 @@ public class LocalBufferPoolTest extends TestLogger {
 
     @Test
     @SuppressWarnings("unchecked")
-    public void testConcurrentRequestRecycle() throws ExecutionException, InterruptedException {
+    void testConcurrentRequestRecycle() throws ExecutionException, InterruptedException {
         int numConcurrentTasks = 128;
         int numBuffersToRequestPerTask = 1024;
 
@@ -410,85 +416,85 @@ public class LocalBufferPoolTest extends TestLogger {
         }
 
         for (int i = 0; i < numConcurrentTasks; i++) {
-            assertTrue(taskResults[i].get());
+            assertThat(taskResults[i].get()).isTrue();
         }
     }
 
     @Test
-    public void testBoundedBuffer() throws Exception {
+    void testBoundedBuffer() {
         localBufferPool.lazyDestroy();
 
         localBufferPool = new LocalBufferPool(networkBufferPool, 1, 2);
-        assertEquals(1, localBufferPool.getNumberOfAvailableMemorySegments());
-        assertEquals(2, localBufferPool.getMaxNumberOfMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isOne();
+        assertThat(localBufferPool.getMaxNumberOfMemorySegments()).isEqualTo(2);
 
         Buffer buffer1, buffer2;
 
         // check min number of buffers:
         localBufferPool.setNumBuffers(1);
-        assertEquals(1, localBufferPool.getNumberOfAvailableMemorySegments());
-        assertNotNull(buffer1 = localBufferPool.requestBuffer());
-        assertEquals(0, localBufferPool.getNumberOfAvailableMemorySegments());
-        assertNull(localBufferPool.requestBuffer());
-        assertEquals(0, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isOne();
+        assertThat(buffer1 = localBufferPool.requestBuffer()).isNotNull();
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isZero();
+        assertThat(localBufferPool.requestBuffer()).isNull();
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isZero();
         buffer1.recycleBuffer();
-        assertEquals(1, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isOne();
 
         // check max number of buffers:
         localBufferPool.setNumBuffers(2);
-        assertEquals(1, localBufferPool.getNumberOfAvailableMemorySegments());
-        assertNotNull(buffer1 = localBufferPool.requestBuffer());
-        assertEquals(1, localBufferPool.getNumberOfAvailableMemorySegments());
-        assertNotNull(buffer2 = localBufferPool.requestBuffer());
-        assertEquals(0, localBufferPool.getNumberOfAvailableMemorySegments());
-        assertNull(localBufferPool.requestBuffer());
-        assertEquals(0, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isOne();
+        assertThat(buffer1 = localBufferPool.requestBuffer()).isNotNull();
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isOne();
+        assertThat(buffer2 = localBufferPool.requestBuffer()).isNotNull();
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isZero();
+        assertThat(localBufferPool.requestBuffer()).isNull();
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isZero();
         buffer1.recycleBuffer();
-        assertEquals(1, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isOne();
         buffer2.recycleBuffer();
-        assertEquals(2, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isEqualTo(2);
 
         // try to set too large buffer size:
         localBufferPool.setNumBuffers(3);
-        assertEquals(2, localBufferPool.getNumberOfAvailableMemorySegments());
-        assertNotNull(buffer1 = localBufferPool.requestBuffer());
-        assertEquals(1, localBufferPool.getNumberOfAvailableMemorySegments());
-        assertNotNull(buffer2 = localBufferPool.requestBuffer());
-        assertEquals(0, localBufferPool.getNumberOfAvailableMemorySegments());
-        assertNull(localBufferPool.requestBuffer());
-        assertEquals(0, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isEqualTo(2);
+        assertThat(buffer1 = localBufferPool.requestBuffer()).isNotNull();
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isOne();
+        assertThat(buffer2 = localBufferPool.requestBuffer()).isNotNull();
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isZero();
+        assertThat(localBufferPool.requestBuffer()).isNull();
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isZero();
         buffer1.recycleBuffer();
-        assertEquals(1, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isOne();
         buffer2.recycleBuffer();
-        assertEquals(2, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isEqualTo(2);
 
         // decrease size again
         localBufferPool.setNumBuffers(1);
-        assertEquals(1, localBufferPool.getNumberOfAvailableMemorySegments());
-        assertNotNull(buffer1 = localBufferPool.requestBuffer());
-        assertEquals(0, localBufferPool.getNumberOfAvailableMemorySegments());
-        assertNull(localBufferPool.requestBuffer());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isOne();
+        assertThat(buffer1 = localBufferPool.requestBuffer()).isNotNull();
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isZero();
+        assertThat(localBufferPool.requestBuffer()).isNull();
         buffer1.recycleBuffer();
-        assertEquals(1, localBufferPool.getNumberOfAvailableMemorySegments());
+        assertThat(localBufferPool.getNumberOfAvailableMemorySegments()).isOne();
     }
 
     /** Moves around availability of a {@link LocalBufferPool} with varying capacity. */
     @Test
-    public void testMaxBuffersPerChannelAndAvailability() throws Exception {
+    void testMaxBuffersPerChannelAndAvailability() throws Exception {
         localBufferPool.lazyDestroy();
         localBufferPool = new LocalBufferPool(networkBufferPool, 1, Integer.MAX_VALUE, 3, 2, 0);
         localBufferPool.setNumBuffers(10);
 
-        assertTrue(localBufferPool.getAvailableFuture().isDone());
+        assertThat(localBufferPool.getAvailableFuture()).isDone();
 
-        // request one segment from subpartitin-0 and subpartition-1 respectively
+        // request one segment from subpartition-0 and subpartition-1 respectively
         final BufferBuilder bufferBuilder01 = localBufferPool.requestBufferBuilderBlocking(0);
         final BufferBuilder bufferBuilder11 = localBufferPool.requestBufferBuilderBlocking(1);
-        assertTrue(localBufferPool.getAvailableFuture().isDone());
+        assertThat(localBufferPool.getAvailableFuture()).isDone();
 
         // request one segment from subpartition-0
         final BufferBuilder bufferBuilder02 = localBufferPool.requestBufferBuilderBlocking(0);
-        assertFalse(localBufferPool.getAvailableFuture().isDone());
+        assertThat(localBufferPool.getAvailableFuture()).isNotDone();
 
         // When this channel reaches maxBuffersPerChannel, LocalBufferPool will be unavailable. But
         // when a memory segment is requested from LocalBufferPool and availableMemorySegments isn't
@@ -497,82 +503,82 @@ public class LocalBufferPoolTest extends TestLogger {
         final BufferBuilder bufferBuilder03 = localBufferPool.requestBufferBuilderBlocking(0);
         final BufferBuilder bufferBuilder21 = localBufferPool.requestBufferBuilderBlocking(2);
         final BufferBuilder bufferBuilder22 = localBufferPool.requestBufferBuilderBlocking(2);
-        assertFalse(localBufferPool.getAvailableFuture().isDone());
+        assertThat(localBufferPool.getAvailableFuture()).isNotDone();
 
         // recycle segments
         bufferBuilder11.close();
-        assertFalse(localBufferPool.getAvailableFuture().isDone());
+        assertThat(localBufferPool.getAvailableFuture()).isNotDone();
         bufferBuilder21.close();
-        assertFalse(localBufferPool.getAvailableFuture().isDone());
+        assertThat(localBufferPool.getAvailableFuture()).isNotDone();
         bufferBuilder02.close();
-        assertFalse(localBufferPool.getAvailableFuture().isDone());
+        assertThat(localBufferPool.getAvailableFuture()).isNotDone();
         bufferBuilder01.close();
-        assertTrue(localBufferPool.getAvailableFuture().isDone());
+        assertThat(localBufferPool.getAvailableFuture()).isDone();
         bufferBuilder03.close();
-        assertTrue(localBufferPool.getAvailableFuture().isDone());
+        assertThat(localBufferPool.getAvailableFuture()).isDone();
         bufferBuilder22.close();
-        assertTrue(localBufferPool.getAvailableFuture().isDone());
+        assertThat(localBufferPool.getAvailableFuture()).isDone();
     }
 
     @Test
-    public void testIsAvailableOrNot() throws InterruptedException {
+    void testIsAvailableOrNot() throws InterruptedException {
 
         // the local buffer pool should be in available state initially
-        assertTrue(localBufferPool.isAvailable());
+        assertThat(localBufferPool.isAvailable()).isTrue();
 
         // request one buffer
         try (BufferBuilder bufferBuilder =
                 checkNotNull(localBufferPool.requestBufferBuilderBlocking())) {
             CompletableFuture<?> availableFuture = localBufferPool.getAvailableFuture();
-            assertFalse(availableFuture.isDone());
+            assertThat(availableFuture).isNotDone();
 
             // set the pool size
             final int numLocalBuffers = 5;
             localBufferPool.setNumBuffers(numLocalBuffers);
-            assertTrue(availableFuture.isDone());
-            assertTrue(localBufferPool.isAvailable());
+            assertThat(availableFuture).isDone();
+            assertThat(localBufferPool.isAvailable()).isTrue();
 
             // drain the local buffer pool
-            final Deque<Buffer> buffers = new ArrayDeque<>(LocalBufferPoolTest.numBuffers);
+            final Deque<Buffer> buffers = new ArrayDeque<>(numBuffers);
             for (int i = 0; i < numLocalBuffers - 1; i++) {
-                assertTrue(localBufferPool.isAvailable());
+                assertThat(localBufferPool.isAvailable()).isTrue();
                 buffers.add(checkNotNull(localBufferPool.requestBuffer()));
             }
-            assertFalse(localBufferPool.isAvailable());
+            assertThat(localBufferPool.isAvailable()).isFalse();
 
             buffers.pop().recycleBuffer();
-            assertTrue(localBufferPool.isAvailable());
+            assertThat(localBufferPool.isAvailable()).isTrue();
 
             // recycle the requested segments to global buffer pool
             for (final Buffer buffer : buffers) {
                 buffer.recycleBuffer();
             }
-            assertTrue(localBufferPool.isAvailable());
+            assertThat(localBufferPool.isAvailable()).isTrue();
 
             // scale down (first buffer still taken), but there should still be one segment locally
             // available
             localBufferPool.setNumBuffers(2);
-            assertTrue(localBufferPool.isAvailable());
+            assertThat(localBufferPool.isAvailable()).isTrue();
 
             final Buffer buffer2 = checkNotNull(localBufferPool.requestBuffer());
-            assertFalse(localBufferPool.isAvailable());
+            assertThat(localBufferPool.isAvailable()).isFalse();
 
             buffer2.recycleBuffer();
-            assertTrue(localBufferPool.isAvailable());
+            assertThat(localBufferPool.isAvailable()).isTrue();
 
             // reset the pool size
             localBufferPool.setNumBuffers(1);
-            assertFalse(localBufferPool.getAvailableFuture().isDone());
+            assertThat(localBufferPool.getAvailableFuture()).isNotDone();
             // recycle the requested buffer
         }
 
-        assertTrue(localBufferPool.isAvailable());
-        assertTrue(localBufferPool.getAvailableFuture().isDone());
+        assertThat(localBufferPool.isAvailable()).isTrue();
+        assertThat(localBufferPool.getAvailableFuture()).isDone();
     }
 
     /** For FLINK-20547: https://issues.apache.org/jira/browse/FLINK-20547. */
     @Test
-    public void testConsistentAvailability() throws Exception {
+    void testConsistentAvailability() throws Exception {
         NetworkBufferPool globalPool = new TestNetworkBufferPool(numBuffers, memorySegmentSize);
         try {
             BufferPool localPool = new LocalBufferPool(globalPool, 1);
@@ -587,7 +593,7 @@ public class LocalBufferPoolTest extends TestLogger {
     }
 
     @Test
-    public void testOverdraftBufferAndAvailability() throws Exception {
+    void testOverdraftBufferAndAvailability() throws Exception {
         for (int maxOverdraftBuffers = 0; maxOverdraftBuffers < 3; maxOverdraftBuffers++) {
             useAllOverdraftBuffersAndCheckIsLegal(4, 3, maxOverdraftBuffers, 2, 1);
             useAllOverdraftBuffersAndCheckIsLegal(4, 3, maxOverdraftBuffers, 2, 2);
@@ -627,7 +633,7 @@ public class LocalBufferPoolTest extends TestLogger {
         for (int i = 0; i < poolSize; i++) {
             int targetChannel = i % availableChannels;
             BufferBuilder bufferBuilder = bufferPool.requestBufferBuilder(targetChannel);
-            assertNotNull(bufferBuilder);
+            assertThat(bufferBuilder).isNotNull();
             closeableRegistryMap
                     .computeIfAbsent(targetChannel, channel -> new AutoCloseableRegistry())
                     .registerCloseable(bufferBuilder);
@@ -641,7 +647,7 @@ public class LocalBufferPoolTest extends TestLogger {
         for (int i = 0; i < maxOverdraftBuffers; i++) {
             int targetChannel = i % availableChannels;
             BufferBuilder bufferBuilder = bufferPool.requestBufferBuilder(targetChannel);
-            assertNotNull(bufferBuilder);
+            assertThat(bufferBuilder).isNotNull();
             overdraftCloseableRegistry.registerCloseable(bufferBuilder);
             int numberOfRequestedOverdraftBuffer = i + 1;
             assertRequestedBufferAndIsAvailable(
@@ -652,7 +658,7 @@ public class LocalBufferPoolTest extends TestLogger {
         }
 
         for (int i = 0; i < numberOfChannels; i++) {
-            assertNull(bufferPool.requestBufferBuilder(i));
+            assertThat(bufferPool.requestBufferBuilder(i)).isNull();
             assertRequestedBufferAndIsAvailable(
                     bufferPool, maxOverdraftBuffers, poolSize + maxOverdraftBuffers, false);
         }
@@ -678,11 +684,11 @@ public class LocalBufferPoolTest extends TestLogger {
         if (numberOfRequestedOverdraftBuffer > 0) {
             checkArgument(!isAvailable);
         }
-        assertEquals(
-                numberOfRequestedOverdraftBuffer,
-                bufferPool.getNumberOfRequestedOverdraftMemorySegments());
-        assertEquals(numberOfRequestedBuffer, bufferPool.bestEffortGetNumOfUsedBuffers());
-        assertEquals(isAvailable, bufferPool.getAvailableFuture().isDone());
+        assertThat(bufferPool.getNumberOfRequestedOverdraftMemorySegments())
+                .isEqualTo(numberOfRequestedOverdraftBuffer);
+
+        assertThat(bufferPool.bestEffortGetNumOfUsedBuffers()).isEqualTo(numberOfRequestedBuffer);
+        assertThat(bufferPool.getAvailableFuture().isDone()).isEqualTo(isAvailable);
     }
 
     // ------------------------------------------------------------------------
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/TestingBufferListener.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/TestingBufferListener.java
new file mode 100644
index 00000000000..98040016778
--- /dev/null
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/TestingBufferListener.java
@@ -0,0 +1,73 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you under the Apache License, Version 2.0 (the
+ * "License"); you may not use this file except in compliance
+ * with the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.flink.runtime.io.network.buffer;
+
+import java.util.function.Function;
+
+/** Mock {@link BufferListener} for testing. */
+public class TestingBufferListener implements BufferListener {
+
+    private final Function<Buffer, Boolean> notifyBufferAvailableFunction;
+
+    private final Runnable notifyBufferDestroyedRunnable;
+
+    private TestingBufferListener(
+            Function<Buffer, Boolean> notifyBufferAvailableFunction,
+            Runnable notifyBufferDestroyedRunnable) {
+        this.notifyBufferAvailableFunction = notifyBufferAvailableFunction;
+        this.notifyBufferDestroyedRunnable = notifyBufferDestroyedRunnable;
+    }
+
+    @Override
+    public boolean notifyBufferAvailable(Buffer buffer) {
+        return notifyBufferAvailableFunction.apply(buffer);
+    }
+
+    @Override
+    public void notifyBufferDestroyed() {
+        notifyBufferDestroyedRunnable.run();
+    }
+
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    /** Builder for {@link TestingBufferListener}. */
+    public static class Builder {
+        private Function<Buffer, Boolean> notifyBufferAvailableFunction = (ignore) -> false;
+
+        private Runnable notifyBufferDestroyedRunnable = () -> {};
+
+        public Builder setNotifyBufferAvailableFunction(
+                Function<Buffer, Boolean> notifyBufferAvailableFunction) {
+            this.notifyBufferAvailableFunction = notifyBufferAvailableFunction;
+            return this;
+        }
+
+        public Builder setNotifyBufferDestroyedRunnable(Runnable notifyBufferDestroyedRunnable) {
+            this.notifyBufferDestroyedRunnable = notifyBufferDestroyedRunnable;
+            return this;
+        }
+
+        public TestingBufferListener build() {
+            return new TestingBufferListener(
+                    notifyBufferAvailableFunction, notifyBufferDestroyedRunnable);
+        }
+    }
+}


[flink] 02/04: [hotfix] Some minor clean-up and typo fix for LocalBufferPool.

Posted by gu...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

guoweijie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit 0a49f9d1e883154e178c55be839e662d2ef213f5
Author: Weijie Guo <re...@163.com>
AuthorDate: Thu Mar 2 23:45:01 2023 +0800

    [hotfix] Some minor clean-up and typo fix for LocalBufferPool.
---
 .../runtime/io/network/buffer/LocalBufferPool.java | 24 +++++++++-------------
 1 file changed, 10 insertions(+), 14 deletions(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
index 214f5c949d9..7e6168ee362 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
@@ -42,7 +42,7 @@ import static org.apache.flink.util.concurrent.FutureUtils.assertNoException;
  * A buffer pool used to manage a number of {@link Buffer} instances from the {@link
  * NetworkBufferPool}.
  *
- * <p>Buffer requests are mediated to the network buffer pool to ensure dead-lock free operation of
+ * <p>Buffer requests are mediated to the network buffer pool to ensure deadlock free operation of
  * the network stack by limiting the number of buffers per local buffer pool. It also implements the
  * default mechanism for buffer recycling, which ensures that every buffer is ultimately returned to
  * the network buffer pool.
@@ -85,8 +85,7 @@ class LocalBufferPool implements BufferPool {
      * org.apache.flink.runtime.io.network.partition.consumer.BufferManager#bufferQueue} via the
      * {@link #registeredListeners} callback.
      */
-    private final ArrayDeque<MemorySegment> availableMemorySegments =
-            new ArrayDeque<MemorySegment>();
+    private final ArrayDeque<MemorySegment> availableMemorySegments = new ArrayDeque<>();
 
     /**
      * Buffer availability listeners, which need to be notified when a Buffer becomes available.
@@ -144,8 +143,7 @@ class LocalBufferPool implements BufferPool {
      * @param networkBufferPool global network buffer pool to get buffers from
      * @param numberOfRequiredMemorySegments minimum number of network buffers
      */
-    LocalBufferPool(NetworkBufferPool networkBufferPool, int numberOfRequiredMemorySegments)
-            throws IOException {
+    LocalBufferPool(NetworkBufferPool networkBufferPool, int numberOfRequiredMemorySegments) {
         this(
                 networkBufferPool,
                 numberOfRequiredMemorySegments,
@@ -166,8 +164,7 @@ class LocalBufferPool implements BufferPool {
     LocalBufferPool(
             NetworkBufferPool networkBufferPool,
             int numberOfRequiredMemorySegments,
-            int maxNumberOfMemorySegments)
-            throws Exception {
+            int maxNumberOfMemorySegments) {
         this(
                 networkBufferPool,
                 numberOfRequiredMemorySegments,
@@ -194,8 +191,7 @@ class LocalBufferPool implements BufferPool {
             int maxNumberOfMemorySegments,
             int numberOfSubpartitions,
             int maxBuffersPerChannel,
-            int maxOverdraftBuffersPerGate)
-            throws IOException {
+            int maxOverdraftBuffersPerGate) {
         checkArgument(
                 numberOfRequiredMemorySegments > 0,
                 "Required number of memory segments (%s) should be larger than 0.",
@@ -542,7 +538,7 @@ class LocalBufferPool implements BufferPool {
             return AvailabilityStatus.UNAVAILABLE_NEED_NOT_REQUESTING_NOTIFICATION;
         }
         boolean needRequestingNotificationOfGlobalPoolAvailable = false;
-        // There aren't availableMemorySegments and we continue to request new memory segment from
+        // There aren't availableMemorySegments, and we continue to request new memory segment from
         // global pool.
         if (!requestMemorySegmentFromGlobal()) {
             // If we can not get a buffer from global pool, we should request from it when it
@@ -605,7 +601,7 @@ class LocalBufferPool implements BufferPool {
             BufferListener listener, MemorySegment segment) {
         // We do not know which locks have been acquired before the recycle() or are needed in the
         // notification and which other threads also access them.
-        // -> call notifyBufferAvailable() outside of the synchronized block to avoid a deadlock
+        // -> call notifyBufferAvailable() outside the synchronized block to avoid a deadlock
         // (FLINK-9676)
         return listener.notifyBufferAvailable(new NetworkBuffer(segment, this));
     }
@@ -669,7 +665,7 @@ class LocalBufferPool implements BufferPool {
             if (isDestroyed) {
                 // FLINK-19964: when two local buffer pools are released concurrently, one of them
                 // gets buffers assigned
-                // make sure that checkAndUpdateAvailability is not called as it would pro-actively
+                // make sure that checkAndUpdateAvailability is not called as it would proactively
                 // acquire one buffer from NetworkBufferPool.
                 return;
             }
@@ -759,8 +755,8 @@ class LocalBufferPool implements BufferPool {
 
     private static class SubpartitionBufferRecycler implements BufferRecycler {
 
-        private int channel;
-        private LocalBufferPool bufferPool;
+        private final int channel;
+        private final LocalBufferPool bufferPool;
 
         SubpartitionBufferRecycler(int channel, LocalBufferPool bufferPool) {
             this.channel = channel;


[flink] 04/04: [FLINK-31293][runtime] LocalBufferPool request overdraft buffer only when no available buffer and pool size is reached.

Posted by gu...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

guoweijie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit fb6caee13710348a9b53284c2cabbdb2e7aa9739
Author: Weijie Guo <re...@163.com>
AuthorDate: Tue Mar 7 01:12:14 2023 +0800

    [FLINK-31293][runtime] LocalBufferPool request overdraft buffer only when no available buffer and pool size is reached.
    
    This closes #22084
---
 .../runtime/io/network/buffer/LocalBufferPool.java |  21 ++-
 .../io/network/buffer/LocalBufferPoolTest.java     | 180 ++++++++++++++++++++-
 2 files changed, 192 insertions(+), 9 deletions(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
index bb14a7d69c5..87b8a11e2aa 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
@@ -51,8 +51,8 @@ import static org.apache.flink.util.concurrent.FutureUtils.assertNoException;
  * will then lazily return the required number of buffers to the {@link NetworkBufferPool} to match
  * its new size.
  *
- * <p>Availability is defined as returning a segment on a subsequent {@link #requestBuffer()}/
- * {@link #requestBufferBuilder()} and heaving a non-blocking {@link
+ * <p>Availability is defined as returning a non-overdraft segment on a subsequent {@link
+ * #requestBuffer()}/ {@link #requestBufferBuilder()} and heaving a non-blocking {@link
  * #requestBufferBuilderBlocking(int)}. In particular,
  *
  * <ul>
@@ -392,14 +392,16 @@ class LocalBufferPool implements BufferPool {
 
     @Nullable
     private MemorySegment requestMemorySegment(int targetChannel) {
-        MemorySegment segment;
+        MemorySegment segment = null;
         synchronized (availableMemorySegments) {
             checkDestroyed();
 
-            if (availableMemorySegments.isEmpty()) {
-                segment = requestOverdraftMemorySegmentFromGlobal();
-            } else {
+            if (!availableMemorySegments.isEmpty()) {
                 segment = availableMemorySegments.poll();
+            } else if (isRequestedSizeReached()) {
+                // Only when the buffer request reaches the upper limit(i.e. current pool size),
+                // requests an overdraft buffer.
+                segment = requestOverdraftMemorySegmentFromGlobal();
             }
 
             if (segment == null) {
@@ -669,6 +671,13 @@ class LocalBufferPool implements BufferPool {
 
             currentPoolSize = Math.min(numBuffers, maxNumberOfMemorySegments);
 
+            // reset overdraft buffers
+            while (numberOfRequestedOverdraftMemorySegments > 0
+                    && numberOfRequestedMemorySegments < currentPoolSize) {
+                numberOfRequestedOverdraftMemorySegments--;
+                numberOfRequestedMemorySegments++;
+            }
+
             returnExcessMemorySegments();
 
             if (isDestroyed) {
diff --git a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
index c677cd4905f..956d55e5948 100644
--- a/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
+++ b/flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPoolTest.java
@@ -38,8 +38,10 @@ import java.util.ArrayDeque;
 import java.util.ArrayList;
 import java.util.Deque;
 import java.util.HashMap;
+import java.util.LinkedList;
 import java.util.List;
 import java.util.Map;
+import java.util.Queue;
 import java.util.concurrent.Callable;
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.ExecutionException;
@@ -249,15 +251,187 @@ class LocalBufferPoolTest {
         }
     }
 
+    @Test
+    void testDecreasePoolSize() throws Exception {
+        final int maxMemorySegments = 10;
+        final int requiredMemorySegments = 4;
+        final int maxOverdraftBuffers = 2;
+        final int largePoolSize = 5;
+        final int smallPoolSize = 4;
+        LocalBufferPool bufferPool =
+                new LocalBufferPool(
+                        networkBufferPool,
+                        requiredMemorySegments,
+                        maxMemorySegments,
+                        0,
+                        Integer.MAX_VALUE,
+                        maxOverdraftBuffers);
+        Queue<MemorySegment> buffers = new LinkedList<>();
+
+        // set a larger pool size.
+        bufferPool.setNumBuffers(largePoolSize);
+        assertThat(bufferPool.getNumBuffers()).isEqualTo(largePoolSize);
+
+        // request all buffer.
+        for (int i = 0; i < largePoolSize; i++) {
+            buffers.add(bufferPool.requestMemorySegmentBlocking());
+        }
+        assertThat(bufferPool.isAvailable()).isFalse();
+
+        // request 1 overdraft buffers.
+        buffers.add(bufferPool.requestMemorySegmentBlocking());
+        assertThat(bufferPool.getNumberOfRequestedOverdraftMemorySegments()).isOne();
+        assertThat(bufferPool.isAvailable()).isFalse();
+
+        // set a small pool size.
+        bufferPool.setNumBuffers(smallPoolSize);
+        assertThat(bufferPool.getNumBuffers()).isEqualTo(smallPoolSize);
+        assertThat(bufferPool.getNumberOfAvailableMemorySegments()).isZero();
+        assertThat(bufferPool.getNumberOfRequestedOverdraftMemorySegments()).isOne();
+        assertThat(bufferPool.isAvailable()).isFalse();
+        buffers.add(bufferPool.requestMemorySegmentBlocking());
+        assertThat(bufferPool.getNumberOfRequestedOverdraftMemorySegments()).isEqualTo(2);
+        assertThat(bufferPool.isAvailable()).isFalse();
+
+        // return all overdraft buffers.
+        bufferPool.recycle(buffers.poll());
+        assertThat(bufferPool.getNumberOfRequestedOverdraftMemorySegments()).isOne();
+        assertThat(bufferPool.isAvailable()).isFalse();
+        bufferPool.recycle(buffers.poll());
+        assertThat(bufferPool.getNumberOfRequestedOverdraftMemorySegments()).isZero();
+        assertThat(bufferPool.isAvailable()).isFalse();
+
+        // return the excess buffer.
+        bufferPool.recycle(buffers.poll());
+        assertThat(bufferPool.isAvailable()).isFalse();
+        // return non-excess buffers.
+        bufferPool.recycle(buffers.poll());
+        assertThat(bufferPool.getNumberOfAvailableMemorySegments()).isOne();
+        assertThat(bufferPool.isAvailable()).isTrue();
+
+        while (!buffers.isEmpty()) {
+            bufferPool.recycle(buffers.poll());
+        }
+        bufferPool.lazyDestroy();
+    }
+
+    @Test
+    void testIncreasePoolSize() throws Exception {
+        final int maxMemorySegments = 100;
+        final int requiredMemorySegments = 5;
+        final int smallPoolSize = 5;
+        final int maxOverdraftBuffers = 2;
+        // test increase pool size not exceed total buffers.
+        testIncreasePoolSizeInternal(
+                maxMemorySegments,
+                requiredMemorySegments,
+                smallPoolSize,
+                6,
+                maxOverdraftBuffers,
+                1,
+                0,
+                false);
+        // test increase pool size to total buffers.
+        testIncreasePoolSizeInternal(
+                maxMemorySegments,
+                requiredMemorySegments,
+                smallPoolSize,
+                7,
+                maxOverdraftBuffers,
+                0,
+                0,
+                false);
+        // test increase pool size exceed total buffers.
+        testIncreasePoolSizeInternal(
+                maxMemorySegments,
+                requiredMemorySegments,
+                smallPoolSize,
+                8,
+                maxOverdraftBuffers,
+                0,
+                1,
+                true);
+        // test increase pool size exceed total buffers and reach maxMemorySegments.
+        testIncreasePoolSizeInternal(10, 8, 8, 10, maxOverdraftBuffers, 0, 0, false);
+    }
+
+    void testIncreasePoolSizeInternal(
+            int maxMemorySegments,
+            int requiredMemorySegments,
+            int smallPoolSize,
+            int largePoolSize,
+            int maxOverdraftBuffers,
+            int numOverdraftBuffersAfterIncreasePoolSize,
+            int numAvailableBuffersAfterIncreasePoolSize,
+            boolean isAvailableAfterIncreasePoolSize)
+            throws Exception {
+        LocalBufferPool bufferPool =
+                new LocalBufferPool(
+                        networkBufferPool,
+                        requiredMemorySegments,
+                        maxMemorySegments,
+                        0,
+                        Integer.MAX_VALUE,
+                        maxOverdraftBuffers);
+        List<MemorySegment> buffers = new ArrayList<>();
+
+        // set a small pool size.
+        bufferPool.setNumBuffers(smallPoolSize);
+        assertThat(bufferPool.getNumBuffers()).isEqualTo(smallPoolSize);
+
+        // request all buffer.
+        for (int i = 0; i < smallPoolSize; i++) {
+            buffers.add(bufferPool.requestMemorySegmentBlocking());
+        }
+        assertThat(bufferPool.isAvailable()).isFalse();
+
+        // request all overdraft buffers.
+        for (int i = 0; i < maxOverdraftBuffers; i++) {
+            buffers.add(bufferPool.requestMemorySegmentBlocking());
+        }
+        assertThat(bufferPool.requestMemorySegment()).isNull();
+        assertThat(bufferPool.getNumberOfRequestedOverdraftMemorySegments())
+                .isEqualTo(maxOverdraftBuffers);
+        assertThat(bufferPool.isAvailable()).isFalse();
+
+        // set a large pool size.
+        bufferPool.setNumBuffers(largePoolSize);
+        assertThat(bufferPool.getNumBuffers()).isEqualTo(largePoolSize);
+        assertThat(bufferPool.getNumberOfAvailableMemorySegments())
+                .isEqualTo(numAvailableBuffersAfterIncreasePoolSize);
+        assertThat(bufferPool.getNumberOfRequestedOverdraftMemorySegments())
+                .isEqualTo(numOverdraftBuffersAfterIncreasePoolSize);
+        assertThat(bufferPool.isAvailable()).isEqualTo(isAvailableAfterIncreasePoolSize);
+
+        for (MemorySegment buffer : buffers) {
+            bufferPool.recycle(buffer);
+        }
+        bufferPool.lazyDestroy();
+    }
+
     @Test
     @Timeout(30)
-    void testRequestBuffersOnRecycle() throws Exception {
-        BufferPool bufferPool1 = networkBufferPool.createBufferPool(512, 2048);
+    void testRequestBufferOnRecycleWithOverdraft() throws Exception {
+        testRequestBuffersOnRecycle(true);
+    }
+
+    @Test
+    @Timeout(30)
+    void testRequestBufferOnRecycleWithoutOverdraft() throws Exception {
+        testRequestBuffersOnRecycle(false);
+    }
+
+    private void testRequestBuffersOnRecycle(boolean supportOverdraftBuffer) throws Exception {
+        BufferPool bufferPool1 =
+                networkBufferPool.createBufferPool(
+                        512, 2048, 0, Integer.MAX_VALUE, supportOverdraftBuffer ? 5 : 0);
         List<MemorySegment> segments = new ArrayList<>();
         for (int i = 0; i < 1023; i++) {
             segments.add(bufferPool1.requestMemorySegmentBlocking());
         }
-        BufferPool bufferPool2 = networkBufferPool.createBufferPool(512, 512);
+        BufferPool bufferPool2 =
+                networkBufferPool.createBufferPool(
+                        512, 512, 0, Integer.MAX_VALUE, supportOverdraftBuffer ? 5 : 0);
         List<MemorySegment> segments2 = new ArrayList<>();
         CheckedThread checkedThread =
                 new CheckedThread() {


[flink] 03/04: [hotfix] Add missing @GuardedBy annotation for LocalBufferPool.

Posted by gu...@apache.org.
This is an automated email from the ASF dual-hosted git repository.

guoweijie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/flink.git

commit cfe393dd87080e8d5c624a385189a95d92db7da9
Author: Weijie Guo <re...@163.com>
AuthorDate: Fri Mar 3 00:03:35 2023 +0800

    [hotfix] Add missing @GuardedBy annotation for LocalBufferPool.
---
 .../flink/runtime/io/network/buffer/LocalBufferPool.java  | 15 ++++++++++++++-
 1 file changed, 14 insertions(+), 1 deletion(-)

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
index 7e6168ee362..bb14a7d69c5 100644
--- a/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
+++ b/flink-runtime/src/main/java/org/apache/flink/runtime/io/network/buffer/LocalBufferPool.java
@@ -295,7 +295,9 @@ class LocalBufferPool implements BufferPool {
 
     @VisibleForTesting
     public int getNumberOfRequestedOverdraftMemorySegments() {
-        return numberOfRequestedOverdraftMemorySegments;
+        synchronized (availableMemorySegments) {
+            return numberOfRequestedOverdraftMemorySegments;
+        }
     }
 
     @Override
@@ -312,6 +314,8 @@ class LocalBufferPool implements BufferPool {
         }
     }
 
+    // suppress the FieldAccessNotGuarded warning as this method is unsafe by design.
+    @SuppressWarnings("FieldAccessNotGuarded")
     @Override
     public int bestEffortGetNumOfUsedBuffers() {
         return Math.max(
@@ -413,6 +417,7 @@ class LocalBufferPool implements BufferPool {
         return segment;
     }
 
+    @GuardedBy("availableMemorySegments")
     private void checkDestroyed() {
         if (isDestroyed) {
             throw new CancelTaskException("Buffer pool has already been destroyed.");
@@ -424,6 +429,7 @@ class LocalBufferPool implements BufferPool {
         return requestMemorySegment(UNKNOWN_CHANNEL);
     }
 
+    @GuardedBy("availableMemorySegments")
     private boolean requestMemorySegmentFromGlobal() {
         assert Thread.holdsLock(availableMemorySegments);
 
@@ -444,6 +450,7 @@ class LocalBufferPool implements BufferPool {
         return false;
     }
 
+    @GuardedBy("availableMemorySegments")
     private MemorySegment requestOverdraftMemorySegmentFromGlobal() {
         assert Thread.holdsLock(availableMemorySegments);
 
@@ -499,6 +506,7 @@ class LocalBufferPool implements BufferPool {
         mayNotifyAvailable(toNotify);
     }
 
+    @GuardedBy("availableMemorySegments")
     private boolean shouldBeAvailable() {
         assert Thread.holdsLock(availableMemorySegments);
 
@@ -551,6 +559,7 @@ class LocalBufferPool implements BufferPool {
                 shouldBeAvailable(), needRequestingNotificationOfGlobalPoolAvailable);
     }
 
+    @GuardedBy("availableMemorySegments")
     private void checkConsistentAvailability() {
         assert Thread.holdsLock(availableMemorySegments);
 
@@ -719,6 +728,7 @@ class LocalBufferPool implements BufferPool {
         }
     }
 
+    @GuardedBy("availableMemorySegments")
     private void returnMemorySegment(MemorySegment segment) {
         assert Thread.holdsLock(availableMemorySegments);
 
@@ -731,6 +741,7 @@ class LocalBufferPool implements BufferPool {
         networkBufferPool.recyclePooledMemorySegment(segment);
     }
 
+    @GuardedBy("availableMemorySegments")
     private void returnExcessMemorySegments() {
         assert Thread.holdsLock(availableMemorySegments);
 
@@ -744,11 +755,13 @@ class LocalBufferPool implements BufferPool {
         }
     }
 
+    @GuardedBy("availableMemorySegments")
     private boolean hasExcessBuffers() {
         return numberOfRequestedOverdraftMemorySegments > 0
                 || numberOfRequestedMemorySegments > currentPoolSize;
     }
 
+    @GuardedBy("availableMemorySegments")
     private boolean isRequestedSizeReached() {
         return numberOfRequestedMemorySegments >= currentPoolSize;
     }