[server] Prevent WAL memory pool deadlocks

fluss-common/src/main/java/org/apache/fluss/memory/LazyMemorySegmentPool.java

@@ -22,6 +22,8 @@
 import org.apache.fluss.config.ConfigOptions;
 import org.apache.fluss.config.Configuration;
 import org.apache.fluss.exception.FlussRuntimeException;
+import org.apache.fluss.exception.RecordTooLargeException;
+import org.apache.fluss.exception.TimeoutException;
 
 import javax.annotation.concurrent.GuardedBy;
 import javax.annotation.concurrent.ThreadSafe;
@@ -33,7 +35,9 @@
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.Deque;
+import java.util.LinkedHashSet;
 import java.util.List;
+import java.util.Set;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Condition;
 import java.util.concurrent.locks.ReentrantLock;
@@ -66,6 +70,14 @@ public class LazyMemorySegmentPool implements MemorySegmentPool, Closeable {
 
     private int pageUsage;
 
+    @GuardedBy("lock")
+    private final Set<Allocation> allocations = new LinkedHashSet<>();
+
+    private final Condition allocationChanged = lock.newCondition();
+
+    @GuardedBy("lock")
+    private int waitingAllocations;
+
     @VisibleForTesting
     LazyMemorySegmentPool(
             int maxPages, int pageSize, long maxTimeToBlockMs, long perRequestMemorySize) {
@@ -231,6 +243,7 @@ public void returnAll(List<MemorySegment> memory) {
                     }
                     pageUsage = newPageUsage;
                     cachePages.addAll(memory);
+                    allocationChanged.signalAll();
                     for (int i = 0; i < memory.size() && !waiters.isEmpty(); i++) {
                         waiters.peekFirst().signal();
                     }
@@ -255,6 +268,7 @@ public void close() {
                     closed = true;
                     cachePages.clear();
                     waiters.forEach(Condition::signal);
+                    allocationChanged.signalAll();
                 });
     }
 
@@ -265,11 +279,144 @@ private void checkClosed() {
     }
 
     public int queued() {
-        return inLock(lock, waiters::size);
+        return inLock(lock, () -> waiters.size() + waitingAllocations);
     }
 
     @VisibleForTesting
     public List<MemorySegment> getAllCachePages() {
         return cachePages;
     }
+
+    @Override
+    public MemoryAllocation newAllocation() {
+        return inLock(
+                lock,
+                () -> {
+                    checkClosed();
+                    Allocation allocation = new Allocation();
+                    allocations.add(allocation);
+                    return allocation;
+                });
+    }
+
+    /** Called only under memory pressure, with the pool lock held. */
+    private void resolveAllocationDeadlock() {
+        int heldPages = 0;
+        Allocation victim = null;
+        for (Allocation allocation : allocations) {
+            int held = allocation.pages.size();
+            heldPages += held;
+            if (held > 0) {
+                // An active owner can still finish, or an aborted owner is already unwinding.
+                if (allocation.pendingPages == 0 || allocation.aborted) {
+                    return;
+                }
+                victim = allocation;
+            }
+            if (allocation.pendingPages > 0 && allocation.pendingPages <= maxPages - pageUsage) {
+                return;
+            }
+        }
+        // Pages outside allocation scopes may be returned independently.
+        if (heldPages == pageUsage && victim != null) {
+            // Registration order keeps older operations alive when holders block each other.
+            victim.aborted = true;
+            allocationChanged.signalAll();
+        }
+    }
+
+    private final class Allocation extends MemoryAllocation {
+        private int pendingPages;
+        private boolean aborted;
+
+        private Allocation() {
+            super(LazyMemorySegmentPool.this);
+        }
+
+        @Override
+        public List<MemorySegment> allocatePages(int required) throws IOException {
+            checkArgument(required > 0, "Requested pages must be positive.");
+            lock.lock();
+            try {
+                checkAllocationOpen();
+                if (required > maxPages - pages.size()) {
+                    aborted = true;
+                    throw new RecordTooLargeException(
+                            "Memory allocation exceeds the memory pool capacity of "
+                                    + totalSize()
+                                    + " bytes: held pages="
+                                    + pages.size()
+                                    + ", requested pages="
+                                    + required
+                                    + ", page size="
+                                    + pageSize);
+                }
+                if (required > maxPages - pageUsage) {
+                    awaitPages(required);
+                }
+                lazilyAllocatePages(required);
+                List<MemorySegment> allocated = drain(required);
+                if (required == 1) {
+                    pages.add(allocated.get(0));
+                } else {
+                    pages.addAll(allocated);
+                }
+                return allocated;
+            } finally {
+                lock.unlock();
+            }
+        }
+
+        private void awaitPages(int required) {
+            pendingPages = required;
+            waitingAllocations++;
+            long remaining = TimeUnit.MILLISECONDS.toNanos(maxTimeToBlockMs);
+            try {
+                while (required > maxPages - pageUsage) {
+                    resolveAllocationDeadlock();
+                    checkAllocationOpen();
+                    if (remaining <= 0) {
+                        throw new TimeoutException("Timed out waiting for memory allocation.");
+                    }
+                    remaining = allocationChanged.awaitNanos(remaining);
+                    checkAllocationOpen();
+                }
+            } catch (InterruptedException e) {
+                Thread.currentThread().interrupt();
+                throw new FlussRuntimeException(e);
+            } finally {
+                pendingPages = 0;
+                waitingAllocations--;
+            }
+        }
+
+        private void checkAllocationOpen() {
+            checkClosed();
+            if (closed) {
+                throw new IllegalStateException("Memory allocation is closed.");
+            }
+            if (aborted) {
+                // Use the existing retryable wire error so older clients can retry as well.
+                throw new TimeoutException(
+                        "Memory allocation aborted because blocked allocations cannot make progress. "
+                                + "Release the allocation and retry the operation.");
+            }
+        }
+
+        @Override
+        public void returnAll(List<MemorySegment> memory) {
+            inLock(lock, () -> super.returnAll(memory));
+        }
+
+        @Override
+        public void close() {
+            inLock(
+                    lock,
+                    () -> {
+                        super.close();
+                        allocations.remove(this);
+                        allocationChanged.signalAll();
+                    });
+        }
+    }
 }

fluss-common/src/main/java/org/apache/fluss/memory/MemoryAllocation.java

@@ -0,0 +1,135 @@
+/*
+ * 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.fluss.memory;
+
+import org.apache.fluss.annotation.Internal;
+
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Collections;
+import java.util.IdentityHashMap;
+import java.util.List;
+import java.util.Set;
+
+import static org.apache.fluss.utils.Preconditions.checkState;
+
+/**
+ * Pages owned by one operation and returned together on close. Use with try-with-resources so
+ * failed or cancelled operations release their pages. An allocation has a single allocating thread;
+ * closing it must not race with that thread or with users of its pages.
+ */
+@Internal
+public class MemoryAllocation implements MemorySegmentPool, AutoCloseable {
+
+    private final MemorySegmentPool pool;
+    protected final List<MemorySegment> pages = new ArrayList<>();
+    protected boolean closed;
+
+    MemoryAllocation(MemorySegmentPool pool) {
+        this.pool = pool;
+    }
+
+    @Override
+    public MemorySegment nextSegment() throws IOException {
+        return allocatePages(1).get(0);
+    }
+
+    @Override
+    public List<MemorySegment> allocatePages(int required) throws IOException {
+        checkState(!closed, "Memory allocation is closed.");
+        List<MemorySegment> allocated = pool.allocatePages(required);
+        pages.addAll(allocated);
+        return allocated;
+    }
+
+    @Override
+    public void returnPage(MemorySegment segment) {
+        returnAll(Collections.singletonList(segment));
+    }
+
+    @Override
+    public void returnAll(List<MemorySegment> memory) {
+        checkState(!closed, "Memory allocation is closed.");
+        if (memory.isEmpty()) {
+            return;
+        }
+        checkState(memory.size() <= pages.size(), "Returned more pages than this allocation owns.");
+        if (ownsAllInOrder(memory)) {
+            pool.returnAll(memory);
+            pages.clear();
+            return;
+        }
+
+        // Validate the entire return before publishing any page to the pool.
+        Set<MemorySegment> returned = Collections.newSetFromMap(new IdentityHashMap<>());
+        for (MemorySegment page : memory) {
+            checkState(returned.add(page), "Page is returned more than once.");
+        }
+        int owned = 0;
+        for (MemorySegment page : pages) {
+            if (returned.contains(page)) {
+                owned++;
+            }
+        }
+        checkState(owned == returned.size(), "Page does not belong to this allocation.");
+        pool.returnAll(memory);
+        pages.removeIf(returned::contains);
+    }
+
+    private boolean ownsAllInOrder(List<MemorySegment> memory) {
+        if (memory.size() != pages.size()) {
+            return false;
+        }
+        int index = 0;
+        for (MemorySegment page : memory) {
+            if (page != pages.get(index++)) {
+                return false;
+            }
+        }
+        return true;
+    }
+
+    @Override
+    public int pageSize() {
+        return pool.pageSize();
+    }
+
+    @Override
+    public long totalSize() {
+        return pool.totalSize();
+    }
+
+    @Override
+    public int freePages() {
+        return pool.freePages();
+    }
+
+    @Override
+    public long availableMemory() {
+        return pool.availableMemory();
+    }
+
+    @Override
+    public void close() {
+        if (!closed) {
+            pool.returnAll(pages);
+            pages.clear();
+            closed = true;
+        }
+    }
+}

fluss-common/src/main/java/org/apache/fluss/memory/MemorySegmentPool.java

@@ -28,6 +28,15 @@
 @Internal
 public interface MemorySegmentPool {
 
+    /**
+     * Opens an allocation whose pages are returned together on close. The caller must release
+     * references to those pages before closing it. Bounded pools may abort a blocked allocation to
+     * allow another allocation to finish; callers must unwind and close the aborted allocation.
+     */
+    default MemoryAllocation newAllocation() {
+        return new MemoryAllocation(this);
+    }
+
     /**
      * Get the page size of each page this pool holds.
      *