feat: add MemorySafeLinkedBlockingQueue

core/common/src/main/java/com/alipay/sofa/rpc/common/threadpool/MemoryLimitCalculator.java

@@ -0,0 +1,80 @@
+/*
+ * 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 com.alipay.sofa.rpc.common.threadpool;
+
+import java.util.concurrent.Executors;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * {@link java.lang.Runtime#freeMemory()} technology is used to calculate the
+ * memory limit by using the percentage of the current maximum available memory,
+ * which can be used with {@link MemoryLimiter}.
+ *
+ * @see MemoryLimiter
+ * @see <a href="https://github.com/apache/incubator-shenyu/blob/master/shenyu-common/src/main/java/org/apache/shenyu/common/concurrent/MemoryLimitCalculator.java">MemoryLimitCalculator</a>
+ */
+public class MemoryLimitCalculator {
+
+    private static volatile long maxAvailable;
+
+    private static final ScheduledExecutorService SCHEDULER = Executors.newSingleThreadScheduledExecutor();
+
+    static {
+        // immediately refresh when this class is loaded to prevent maxAvailable from being 0
+        refresh();
+        // check every 50 ms to improve performance
+        SCHEDULER.scheduleWithFixedDelay(MemoryLimitCalculator::refresh, 50, 50, TimeUnit.MILLISECONDS);
+        Runtime.getRuntime().addShutdownHook(new Thread(SCHEDULER::shutdown));
+    }
+
+    private static void refresh() {
+        maxAvailable = Runtime.getRuntime().freeMemory();
+    }
+
+    /**
+     * Get the maximum available memory of the current JVM.
+     *
+     * @return maximum available memory
+     */
+    public static long maxAvailable() {
+        return maxAvailable;
+    }
+
+    /**
+     * Take the current JVM's maximum available memory
+     * as a percentage of the result as the limit.
+     *
+     * @param percentage percentage
+     * @return available memory
+     */
+    public static long calculate(final float percentage) {
+        if (percentage <= 0 || percentage > 1) {
+            throw new IllegalArgumentException();
+        }
+        return (long) (maxAvailable() * percentage);
+    }
+
+    /**
+     * By default, it takes 80% of the maximum available memory of the current JVM.
+     *
+     * @return available memory
+     */
+    public static long defaultLimit() {
+        return (long) (maxAvailable() * 0.8);
+    }
+}

core/common/src/main/java/com/alipay/sofa/rpc/common/threadpool/MemoryLimiter.java

@@ -0,0 +1,269 @@
+/*
+ * 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 com.alipay.sofa.rpc.common.threadpool;
+
+import java.lang.instrument.Instrumentation;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.LongAdder;
+import java.util.concurrent.locks.Condition;
+import java.util.concurrent.locks.ReentrantLock;
+
+/**
+ * memory limiter.
+ */
+public class MemoryLimiter {
+
+    private final Instrumentation inst;
+
+    private long                  memoryLimit;
+
+    private final LongAdder       memory      = new LongAdder();
+
+    private final ReentrantLock   acquireLock = new ReentrantLock();
+
+    private final Condition       notLimited  = acquireLock.newCondition();
+
+    private final ReentrantLock   releaseLock = new ReentrantLock();
+
+    private final Condition       notEmpty    = releaseLock.newCondition();
+
+    public MemoryLimiter(Instrumentation inst) {
+        this(Integer.MAX_VALUE, inst);
+    }
+
+    public MemoryLimiter(long memoryLimit, Instrumentation inst) {
+        if (memoryLimit <= 0) {
+            throw new IllegalArgumentException();
+        }
+        this.memoryLimit = memoryLimit;
+        this.inst = inst;
+    }
+
+    public void setMemoryLimit(long memoryLimit) {
+        if (memoryLimit <= 0) {
+            throw new IllegalArgumentException();
+        }
+        this.memoryLimit = memoryLimit;
+    }
+
+    public long getMemoryLimit() {
+        return memoryLimit;
+    }
+
+    public long getCurrentMemory() {
+        return memory.sum();
+    }
+
+    public long getCurrentRemainMemory() {
+        return getMemoryLimit() - getCurrentMemory();
+    }
+
+    private void signalNotEmpty() {
+        releaseLock.lock();
+        try {
+            notEmpty.signal();
+        } finally {
+            releaseLock.unlock();
+        }
+    }
+
+    private void signalNotLimited() {
+        acquireLock.lock();
+        try {
+            notLimited.signal();
+        } finally {
+            acquireLock.unlock();
+        }
+    }
+
+    /**
+     * Locks to prevent both acquires and releases.
+     */
+    private void fullyLock() {
+        acquireLock.lock();
+        releaseLock.lock();
+    }
+
+    /**
+     * Unlocks to allow both acquires and releases.
+     */
+    private void fullyUnlock() {
+        releaseLock.unlock();
+        acquireLock.unlock();
+    }
+
+    public boolean acquire(Object e) {
+        if (e == null) {
+            throw new NullPointerException();
+        }
+        if (memory.sum() >= memoryLimit) {
+            return false;
+        }
+        acquireLock.lock();
+        try {
+            final long sum = memory.sum();
+            final long objectSize = inst.getObjectSize(e);
+            if (sum + objectSize >= memoryLimit) {
+                return false;
+            }
+            memory.add(objectSize);
+            // see https://github.com/apache/incubator-shenyu/pull/3356
+            if (memory.sum() < memoryLimit) {
+                notLimited.signal();
+            }
+        } finally {
+            acquireLock.unlock();
+        }
+        if (memory.sum() > 0) {
+            signalNotEmpty();
+        }
+        return true;
+    }
+
+    public void acquireInterruptibly(Object e) throws InterruptedException {
+        if (e == null) {
+            throw new NullPointerException();
+        }
+        acquireLock.lockInterruptibly();
+        try {
+            final long objectSize = inst.getObjectSize(e);
+            // see https://github.com/apache/incubator-shenyu/pull/3335
+            while (memory.sum() + objectSize >= memoryLimit) {
+                notLimited.await();
+            }
+            memory.add(objectSize);
+            if (memory.sum() < memoryLimit) {
+                notLimited.signal();
+            }
+        } finally {
+            acquireLock.unlock();
+        }
+        if (memory.sum() > 0) {
+            signalNotEmpty();
+        }
+    }
+
+    public boolean acquire(Object e, long timeout, TimeUnit unit) throws InterruptedException {
+        if (e == null) {
+            throw new NullPointerException();
+        }
+        long nanos = unit.toNanos(timeout);
+        acquireLock.lockInterruptibly();
+        try {
+            final long objectSize = inst.getObjectSize(e);
+            while (memory.sum() + objectSize >= memoryLimit) {
+                if (nanos <= 0) {
+                    return false;
+                }
+                nanos = notLimited.awaitNanos(nanos);
+            }
+            memory.add(objectSize);
+            if (memory.sum() < memoryLimit) {
+                notLimited.signal();
+            }
+        } finally {
+            acquireLock.unlock();
+        }
+        if (memory.sum() > 0) {
+            signalNotEmpty();
+        }
+        return true;
+    }
+
+    public void release(Object e) {
+        if (null == e) {
+            return;
+        }
+        if (memory.sum() == 0) {
+            return;
+        }
+        releaseLock.lock();
+        try {
+            final long objectSize = inst.getObjectSize(e);
+            if (memory.sum() > 0) {
+                memory.add(-objectSize);
+                if (memory.sum() > 0) {
+                    notEmpty.signal();
+                }
+            }
+        } finally {
+            releaseLock.unlock();
+        }
+        if (memory.sum() < memoryLimit) {
+            signalNotLimited();
+        }
+    }
+
+    public void releaseInterruptibly(Object e) throws InterruptedException {
+        if (null == e) {
+            return;
+        }
+        releaseLock.lockInterruptibly();
+        try {
+            final long objectSize = inst.getObjectSize(e);
+            while (memory.sum() == 0) {
+                notEmpty.await();
+            }
+            memory.add(-objectSize);
+            if (memory.sum() > 0) {
+                notEmpty.signal();
+            }
+        } finally {
+            releaseLock.unlock();
+        }
+        if (memory.sum() < memoryLimit) {
+            signalNotLimited();
+        }
+    }
+
+    public void releaseInterruptibly(Object e, long timeout, TimeUnit unit) throws InterruptedException {
+        if (null == e) {
+            return;
+        }
+        long nanos = unit.toNanos(timeout);
+        releaseLock.lockInterruptibly();
+        try {
+            final long objectSize = inst.getObjectSize(e);
+            while (memory.sum() == 0) {
+                if (nanos <= 0) {
+                    return;
+                }
+                nanos = notEmpty.awaitNanos(nanos);
+            }
+            memory.add(-objectSize);
+            if (memory.sum() > 0) {
+                notEmpty.signal();
+            }
+        } finally {
+            releaseLock.unlock();
+        }
+        if (memory.sum() < memoryLimit) {
+            signalNotLimited();
+        }
+    }
+
+    public void clear() {
+        fullyLock();
+        try {
+            if (memory.sumThenReset() < memoryLimit) {
+                notLimited.signal();
+            }
+        } finally {
+            fullyUnlock();
+        }
+    }
+}