[bugfix] do not remove from waiters before segment allocated (#974)

Author: zcoo

fluss-common/src/main/java/com/alibaba/fluss/memory/LazyMemorySegmentPool.java

@@ -231,7 +231,7 @@ public void returnAll(List<MemorySegment> memory) {
                     pageUsage = newPageUsage;
                     cachePages.addAll(memory);
                     for (int i = 0; i < memory.size() && !waiters.isEmpty(); i++) {
-                        waiters.pollFirst().signal();
+                        waiters.peekFirst().signal();
                     }
                 });
     }

fluss-common/src/test/java/com/alibaba/fluss/memory/LazyMemorySegmentPoolTest.java

@@ -193,6 +193,30 @@ void testCannotAllocateMorePagesThanAvailable() throws Exception {
                 .hasMessageContaining("Total pages: 2. Requested pages: 3");
     }
 
+    @Test
+    void testComplexDelayedAllocation() throws Exception {
+        LazyMemorySegmentPool pool = buildLazyMemorySegmentSource(8, 1024, Long.MAX_VALUE, 1024);
+        // allocate all segments
+        List<MemorySegment> segments1 = pool.allocatePages(4);
+        List<MemorySegment> segments2 = pool.allocatePages(4);
+
+        CountDownLatch doDealloc1 = asyncReturnAll(pool, segments1);
+        CountDownLatch doDealloc2 = asyncReturnAll(pool, segments2);
+
+        // should block until all memory is returned
+        CountDownLatch allocation = asyncAllocatePages(pool, 8);
+        assertThat(allocation.getCount()).isEqualTo(1);
+        // return a part of memory
+        doDealloc1.countDown();
+
+        // no enough memory allocate
+        assertThat(allocation.await(1, TimeUnit.SECONDS)).isFalse();
+
+        doDealloc2.countDown();
+        // have enough memory to allocate
+        assertThat(allocation.await(1, TimeUnit.SECONDS)).isTrue();
+    }
+
     @Test
     void testDelayedAllocation() throws Exception {
         LazyMemorySegmentPool pool = buildLazyMemorySegmentSource(5, 1024, Long.MAX_VALUE, 1024);