[ISSUE #3763] support custom reject in MemorySafeLinkedBlockingQueue (#3764)

* [ISSUE #3763] support custom reject in MemorySafeLinkedBlockingQueue

* fix code style

Author: loongs-zhang

shenyu-common/src/main/java/org/apache/shenyu/common/concurrent/AbortPolicy.java

@@ -0,0 +1,30 @@
+/*
+ * 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.shenyu.common.concurrent;
+
+/**
+ * A handler for rejected element that throws a
+ * {@code RejectException}.
+ */
+public class AbortPolicy<E> implements Rejector<E> {
+
+    @Override
+    public void reject(final E e, final MemorySafeLinkedBlockingQueue<E> queue) {
+        throw new RejectException("no more memory can be used !");
+    }
+}

shenyu-common/src/main/java/org/apache/shenyu/common/concurrent/DiscardOldestPolicy.java

@@ -0,0 +1,30 @@
+/*
+ * 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.shenyu.common.concurrent;
+
+/**
+ * A handler for rejected element that discards the oldest element.
+ */
+public class DiscardOldestPolicy<E> implements Rejector<E> {
+
+    @Override
+    public void reject(final E e, final MemorySafeLinkedBlockingQueue<E> queue) {
+        queue.poll();
+        queue.offer(e);
+    }
+}

shenyu-common/src/main/java/org/apache/shenyu/common/concurrent/DiscardPolicy.java

@@ -0,0 +1,30 @@
+/*
+ * 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.shenyu.common.concurrent;
+
+/**
+ * A handler for rejected element that silently discards the
+ * rejected element.
+ */
+public class DiscardPolicy<E> implements Rejector<E> {
+
+    @Override
+    public void reject(final E e, final MemorySafeLinkedBlockingQueue<E> queue) {
+
+    }
+}

shenyu-common/src/main/java/org/apache/shenyu/common/concurrent/MemorySafeLinkedBlockingQueue.java

@@ -32,15 +32,21 @@ public class MemorySafeLinkedBlockingQueue<E> extends LinkedBlockingQueue<E> {
 
     private int maxFreeMemory;
 
+    private Rejector<E> rejector;
+
     public MemorySafeLinkedBlockingQueue(final int maxFreeMemory) {
         super(Integer.MAX_VALUE);
         this.maxFreeMemory = maxFreeMemory;
+        //default as DiscardPolicy to ensure compatibility with the old version
+        this.rejector = new DiscardPolicy<>();
     }
 
     public MemorySafeLinkedBlockingQueue(final Collection<? extends E> c,
                                          final int maxFreeMemory) {
         super(c);
         this.maxFreeMemory = maxFreeMemory;
+        //default as DiscardPolicy to ensure compatibility with the old version
+        this.rejector = new DiscardPolicy<>();
     }
 
     /**
@@ -61,6 +67,15 @@ public int getMaxFreeMemory() {
         return maxFreeMemory;
     }
 
+    /**
+     * set the rejector.
+     *
+     * @param rejector the rejector
+     */
+    public void setRejector(final Rejector<E> rejector) {
+        this.rejector = rejector;
+    }
+
     /**
      * determine if there is any remaining free memory.
      *
@@ -75,15 +90,24 @@ public void put(final E e) throws InterruptedException {
         if (hasRemainedMemory()) {
             super.put(e);
         }
+        rejector.reject(e, this);
     }
 
     @Override
     public boolean offer(final E e, final long timeout, final TimeUnit unit) throws InterruptedException {
-        return hasRemainedMemory() && super.offer(e, timeout, unit);
+        if (!hasRemainedMemory()) {
+            rejector.reject(e, this);
+            return false;
+        }
+        return super.offer(e, timeout, unit);
     }
 
     @Override
     public boolean offer(final E e) {
-        return hasRemainedMemory() && super.offer(e);
+        if (!hasRemainedMemory()) {
+            rejector.reject(e, this);
+            return false;
+        }
+        return super.offer(e);
     }
 }

shenyu-common/src/main/java/org/apache/shenyu/common/concurrent/RejectException.java

@@ -0,0 +1,72 @@
+/*
+ * 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.shenyu.common.concurrent;
+
+/**
+ * Exception thrown by an {@link MemorySafeLinkedBlockingQueue}
+ * when a element cannot be accepted.
+ */
+public class RejectException extends RuntimeException {
+
+    private static final long serialVersionUID = -3240015871717170195L;
+
+    /**
+     * Constructs a {@code RejectException} with no detail message.
+     * The cause is not initialized, and may subsequently be
+     * initialized by a call to {@link #initCause(Throwable) initCause}.
+     */
+    public RejectException() {
+    }
+
+    /**
+     * Constructs a {@code RejectException} with the
+     * specified detail message. The cause is not initialized, and may
+     * subsequently be initialized by a call to {@link
+     * #initCause(Throwable) initCause}.
+     *
+     * @param message the detail message
+     */
+    public RejectException(final String message) {
+        super(message);
+    }
+
+    /**
+     * Constructs a {@code RejectException} with the
+     * specified detail message and cause.
+     *
+     * @param message the detail message
+     * @param cause   the cause (which is saved for later retrieval by the
+     *                {@link #getCause()} method)
+     */
+    public RejectException(final String message, final Throwable cause) {
+        super(message, cause);
+    }
+
+    /**
+     * Constructs a {@code RejectException} with the
+     * specified cause.  The detail message is set to {@code (cause ==
+     * null ? null : cause.toString())} (which typically contains
+     * the class and detail message of {@code cause}).
+     *
+     * @param cause the cause (which is saved for later retrieval by the
+     *              {@link #getCause()} method)
+     */
+    public RejectException(final Throwable cause) {
+        super(cause);
+    }
+}

shenyu-common/src/main/java/org/apache/shenyu/common/concurrent/Rejector.java

@@ -0,0 +1,43 @@
+/*
+ * 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.shenyu.common.concurrent;
+
+/**
+ * RejectHandler, it works when you need to custom reject action in
+ * {@link org.apache.shenyu.common.concurrent.MemorySafeLinkedBlockingQueue}.
+ *
+ * @see AbortPolicy
+ * @see DiscardPolicy
+ * @see DiscardOldestPolicy
+ */
+public interface Rejector<E> {
+
+    /**
+     * Method that may be invoked by a {@link MemorySafeLinkedBlockingQueue} when
+     * {@link MemorySafeLinkedBlockingQueue#hasRemainedMemory} return true.
+     * This may occur when no more memory are available because their bounds would be exceeded.
+     *
+     * <p>In the absence of other alternatives, the method may throw an unchecked
+     * {@link RejectException}, which will be propagated to the caller.
+     *
+     * @param e     the element requested to be added
+     * @param queue the queue attempting to add this element
+     * @throws RejectException if there is no more memory
+     */
+    void reject(E e, MemorySafeLinkedBlockingQueue<E> queue);
+}

shenyu-common/src/test/java/org/apache/shenyu/common/concurrent/MemorySafeLinkedBlockingQueueTest.java

@@ -21,6 +21,7 @@
 
 import static org.hamcrest.MatcherAssert.assertThat;
 import static org.hamcrest.Matchers.is;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 public class MemorySafeLinkedBlockingQueueTest {
     @Test
@@ -35,4 +36,12 @@ public void test() throws Exception {
         assertThat(queue.offer(() -> {
         }), is(true));
     }
+
+    @Test
+    public void testCustomReject() throws Exception {
+        MemorySafeLinkedBlockingQueue<Runnable> queue = new MemorySafeLinkedBlockingQueue<>(Integer.MAX_VALUE);
+        queue.setRejector(new AbortPolicy<>());
+        assertThrows(RejectException.class, () -> queue.offer(() -> {
+        }));
+    }
 }