zhijiangW commented on a change in pull request #9993: 
[FLINK-14498][runtime]Introduce NetworkBufferPool#isAvailable() for interacting 
with LocalBufferPool.
URL: https://github.com/apache/flink/pull/9993#discussion_r339408191
 
 

 ##########
 File path: 
flink-runtime/src/test/java/org/apache/flink/runtime/io/network/buffer/NetworkBufferPoolTest.java
 ##########
 @@ -518,4 +528,176 @@ public void go() throws Exception {
                        globalPool.destroy();
                }
        }
+
+       /**
+        * Tests {@link NetworkBufferPool#isAvailable()}, verifying that the 
buffer availability is
+        * correctly maintained and the future callback is correctly processed.
+        */
+       @Test
+       public void testBufferAvailabilityAndFutureCallback() throws Exception {
+               final int numBuffers = 10;
+               final int numberOfSegmentsToRequest = 5;
+               final Duration requestSegmentsTimeout = Duration.ofSeconds(10L);
+
+               final NetworkBufferPool globalPool = new NetworkBufferPool(
+                               numBuffers,
+                               128,
+                               numberOfSegmentsToRequest,
+                               requestSegmentsTimeout);
+
+               try {
+                       assertTrue(globalPool.isAvailable().isDone());
+
+                       List<MemorySegment> segments = new 
ArrayList<>(numberOfSegmentsToRequest);
+                       for (int i = 0; i < numberOfSegmentsToRequest; ++i) {
+                               MemorySegment segment = 
globalPool.requestMemorySegment();
+                               assertNotNull(segment);
+                               segments.add(segment);
+                               assertTrue(globalPool.isAvailable().isDone());
+                       }
+
+                       final List<MemorySegment> exclusiveSegments = 
globalPool.requestMemorySegments();
+                       assertEquals(numberOfSegmentsToRequest, 
exclusiveSegments.size());
+
+                       assertFalse(globalPool.isAvailable().isDone());
+                       assertEquals(0, 
globalPool.getNumberOfAvailableMemorySegments());
+
+                       final AtomicBoolean completeFlag = new 
AtomicBoolean(false);
+                       CompletableFuture<?> availableFuture = 
globalPool.isAvailable();
+                       availableFuture.whenComplete((ignored, throwable) -> 
completeFlag.set(true));
+
+                       // recycle one buffer
+                       globalPool.recycle(segments.get(0));
+                       assertTrue(completeFlag.get());
+                       assertTrue(availableFuture.isDone());
+                       assertTrue(globalPool.isAvailable().isDone());
+                       assertEquals(1, 
globalPool.getNumberOfAvailableMemorySegments());
+
+                       CheckedThread asyncRequest = new CheckedThread() {
+                               @Override
+                               public void go() throws Exception {
+                                       
exclusiveSegments.addAll(globalPool.requestMemorySegments());
+                               }
+                       };
+                       asyncRequest.start();
+
+                       // wait until no buffer is available
+                       final Deadline deadline = 
Deadline.fromNow(Duration.ofSeconds(10L));
+                       while (globalPool.getNumberOfAvailableMemorySegments() 
> 0) {
+                               Thread.sleep(50);
+                               if (!deadline.hasTimeLeft()) {
+                                       fail("Waiting timeout.");
+                               }
+                       }
+                       assertFalse(globalPool.isAvailable().isDone());
+                       assertEquals(numberOfSegmentsToRequest, 
exclusiveSegments.size());
+
+                       for (int i = 1; i < numberOfSegmentsToRequest; ++i) {
+                               globalPool.recycle(segments.get(i));
+                       }
+                       segments.clear();
+
+                       asyncRequest.sync();
+                       assertFalse(globalPool.isAvailable().isDone());
+                       assertEquals(numBuffers, exclusiveSegments.size());
 
 Review comment:
   It is not very strict to use `numBuffers` here. Actually it should be `2 * 
numberOfSegmentsToRequest` and it is coincidental that `numBuffers = 2 *  
numberOfSegmentsToRequest`.

----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
 
For queries about this service, please contact Infrastructure at:
us...@infra.apache.org


With regards,
Apache Git Services

Reply via email to