Repository: AutoMQ/automq (Java)

Issue:
[Enhancement] Add quota for low-priority controller mutations
### Who is this for and what problem do they have today?

A large number of compaction or delete object requests may overload the controller event loop.
We should add a mechanism to protect the controller by stopping low-priority requests and allowing it to prioritize high-priority requests, such as broker heartbeat.

### Why is solving this problem impactful?

The high-priority request can not fail; otherwise, the cluster may refuse to serve.



Discussion: Hello 👋 
This issue looks interesting for me and I'd like to take it over, but I'd love to ask few questions upfront 

1. Do I understand it right that the idea, simply saying, is substituting `LinkedBlockingQueue` in `ControllerEventManager` to `PriorityBlockingQueue` ?
2. If I understand this correctly, I see that in [Kafka](https://github.com/apache/kafka/blob/4d3e366bc2487ca0c2321e6e6d7786fb6ed0efa5/core/src/main/scala/kafka/controller/ControllerEventManager.scala#L83) upstream they still use `LinkedBlockingQueue`. Might it be a case that they've already considered this idea and rejected it for any reasons? 
3. Do you have already any statistics about different events and their possibilities to block the event loop? It might be helpful for identifying the types of events that should be with the low priority.

Looking forward for your answer 🙏
> Hello 👋 This issue looks interesting for me and I'd like to take it over, but I'd love to ask few questions upfront
> 
> 1. Do I understand it right that the idea, simply saying, is substituting `LinkedBlockingQueue` in `ControllerEventManager` to `PriorityBlockingQueue` ?
> 2. If I understand this correctly, I see that in [Kafka](https://github.com/apache/kafka/blob/4d3e366bc2487ca0c2321e6e6d7786fb6ed0efa5/core/src/main/scala/kafka/controller/ControllerEventManager.scala#L83) upstream they still use `LinkedBlockingQueue`. Might it be a case that they've already considered this idea and rejected it for any reasons?
> 3. Do you have already any statistics about different events and their possibilities to block the event loop? It might be helpful for identifying the types of events that should be with the low priority.
> 
> Looking forward for your answer 🙏

Using the PriorityBlockingQueue instead of LinkedBlockingQueue will change the order of the requests, potentially causing unexpected problems. One possible solution is to set a time limit for processing low-priority requests. Currently, we rely on a log similar to the one below to identify these low-priority requests:

```
Controller took 12415025 µs to process write event: checkS3ObjectsLifecycle
```
 
Thank you for the clarifications 🙏 Initially I was looking at the wrong direction. I found [that docs](https://docs.automq.com/automq/what-is-automq/compatibility-with-apache-kafka#apache-kafka-compatibility-overview) where was said that AutoMQ doesn't even use old ZK mode and instead AutoMQ always runs in KRaft mode
Object deletion task has been optimized through https://github.com/AutoMQ/automq/pull/1603.

To further prevent Controller Overload (Huge tasks, Long Running task, GC) from causing the entire Broker cluster to be Fenced, `BrokerHeartbeatManager#findOneStaleBroker` adds the following mechanism:
- Add states OVERLOAD_CLOSE, OVERLOAD_OPEN, OVERLOAD_HALF_OPEN
- OVERLOAD_OPEN:
  - If all unfenced brokers timeout, mark the Controller as OVERLOAD_OPEN;
  - Return None;
- OVERLOAD_HALF_OPEN:
  - If in the OVERLOAD_OPEN state, and a new broker HeartBeat arrives, switch to OVERLOAD_HALF_OPEN state;
  - Return None
- OVERLOAD_CLOSE:
  - If OVERLOAD_HALF_OPEN state persists for 30s without deteriorating into OVERLOAD_OPEN, switch to OVERLOAD_CLOSE state;
  - Return the Broker with heartbeat timeout




Fix (patch):
diff --git a/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java b/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java
index 6cb7f615cc..7ce3af55b1 100644
--- a/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java
+++ b/metadata/src/main/java/org/apache/kafka/controller/BrokerHeartbeatManager.java
@@ -21,6 +21,7 @@
 import org.apache.kafka.common.message.BrokerHeartbeatRequestData;
 import org.apache.kafka.common.utils.LogContext;
 import org.apache.kafka.common.utils.Time;
+import org.apache.kafka.controller.stream.OverloadCircuitBreaker;
 import org.apache.kafka.metadata.placement.UsableBroker;
 import org.slf4j.Logger;
 
@@ -162,6 +163,11 @@ BrokerHeartbeatState first() {
             return result == head ? null : result;
         }
 
+        BrokerHeartbeatState last() {
+            BrokerHeartbeatState result = head.prev;
+            return result == head ? null : result;
+        }
+
         /**
          * Add the broker to the list. We start looking for a place to put it at the end
          * of the list.
@@ -251,6 +257,10 @@ public BrokerHeartbeatState next() {
      */
     private final TreeSet<BrokerHeartbeatState> active;
 
+    // AutoMQ inject start
+    private final OverloadCircuitBreaker overloadCircuitBreaker;
+    // AutoMQ inject end
+
     BrokerHeartbeatManager(LogContext logContext,
                            Time time,
                            long sessionTimeoutNs) {
@@ -260,6 +270,10 @@ public BrokerHeartbeatState next() {
         this.brokers = new HashMap<>();
         this.unfenced = new BrokerHeartbeatStateList();
         this.active = new TreeSet<>(MetadataOffsetComparator.INSTANCE);
+
+        // AutoMQ inject start
+        this.overloadCircuitBreaker = new OverloadCircuitBreaker(time);
+        // AutoMQ inject end
     }
 
     // VisibleForTesting
@@ -398,6 +412,9 @@ void touch(int brokerId, boolean fenced, long metadataOffset) {
             if (!broker.shuttingDown()) {
                 active.add(broker);
             }
+            // AutoMQ inject start
+            overloadCircuitBreaker.success();
+            // AutoMQ inject end
         }
     }
 
@@ -434,7 +451,7 @@ void maybeUpdateControlledShutdownOffset(int brokerId, long controlledShutDownOf
      * Return the time in monotonic nanoseconds at which we should check if a broker
      * session needs to be expired.
      */
-    long nextCheckTimeNs() {
+    long nextCheckTimeNs0() {
         BrokerHeartbeatState broker = unfenced.first();
         if (broker == null) {
             return Long.MAX_VALUE;
@@ -449,7 +466,8 @@ long nextCheckTimeNs() {
      *
      * @return      An Optional broker node id.
      */
-    Optional<Integer> findOneStaleBroker() {
+    Optional<Integer> findOneStaleBroker0() {
+
         BrokerHeartbeatStateIterator iterator = unfenced.iterator();
         if (iterator.hasNext()) {
             BrokerHeartbeatState broker = iterator.next();
@@ -469,6 +487,44 @@ Iterator<UsableBroker> usableBrokers(
             idToRack);
     }
 
+    // AutoMQ inject start
+    long nextCheckTimeNs() {
+        if (overloadCircuitBreaker.isOverload()) {
+            return Long.MAX_VALUE;
+        } else {
+            return nextCheckTimeNs0();
+        }
+    }
+
+    Optional<Integer> findOneStaleBroker() {
+        // If the controller is in overload state, we should not fence any broker.
+        if (overloadCircuitBreaker.isOverload()) {
+            return Optional.empty();
+        }
+        Optional<Integer> staleBroker = findOneStaleBroker0();
+        if (staleBroker.isEmpty()) {
+            return staleBroker;
+        }
+        // Check whether the controller is in overload state.
+        if (checkOverload()) {
+            return Optional.empty();
+        } else {
+            return staleBroker;
+        }
+    }
+
+    boolean checkOverload() {
+        // If all unfenced brokers are stale, the controller is in overload state.
+        BrokerHeartbeatState last = unfenced.last();
+        if (last != null && !hasValidSession(last.id)) {
+            // all brokers are stale
+            overloadCircuitBreaker.overload();
+            return true;
+        }
+        return false;
+    }
+    // AutoMQ inject end
+
     static class UsableBrokerIterator implements Iterator<UsableBroker> {
         private final Iterator<BrokerHeartbeatState> iterator;
         private final Function<Integer, Optional<String>> idToRack;
diff --git a/metadata/src/main/java/org/apache/kafka/controller/stream/OverloadCircuitBreaker.java b/metadata/src/main/java/org/apache/kafka/controller/stream/OverloadCircuitBreaker.java
new file mode 100644
index 0000000000..84e7c6543c
--- /dev/null
+++ b/metadata/src/main/java/org/apache/kafka/controller/stream/OverloadCircuitBreaker.java
@@ -0,0 +1,64 @@
+/*
+ * Copyright 2024, AutoMQ HK Limited.
+ *
+ * Use of this software is governed by the Business Source License
+ * included in the file BSL.md
+ *
+ * As of the Change Date specified in that file, in accordance with
+ * the Business Source License, use of this software will be governed
+ * by the Apache License, Version 2.0
+ */
+
+package org.apache.kafka.controller.stream;
+
+import java.util.concurrent.TimeUnit;
+import org.apache.kafka.common.utils.Time;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+public class OverloadCircuitBreaker {
+    private static final Logger LOGGER = LoggerFactory.getLogger(OverloadCircuitBreaker.class);
+    static final long HALF_OPEN_WINDOW_MS = TimeUnit.SECONDS.toMillis(30);
+    private State state = State.OPEN;
+    private long timestamp;
+
+    private final Time time;
+
+    public OverloadCircuitBreaker(Time time) {
+        this.time = time;
+    }
+
+    public void overload() {
+        this.state = State.CLOSED;
+        this.timestamp = time.milliseconds();
+        LOGGER.error("[WARN] The controller is overload, enter overload protecting state.");
+    }
+
+    public void success() {
+        if (this.state == State.OPEN) {
+            return;
+        }
+        long now = time.milliseconds();
+        if (this.state == State.CLOSED) {
+            this.state = State.HALF_OPEN;
+            timestamp = now;
+            LOGGER.info("The controller is recovering from overload, enter half-open state.");
+        } else if (this.state == State.HALF_OPEN) {
+            if (now - timestamp >= HALF_OPEN_WINDOW_MS) {
+                this.state = State.OPEN;
+                LOGGER.info("The controller has recovered from overload, enter open state.");
+            }
+        }
+    }
+
+    public boolean isOverload() {
+        return state != State.OPEN;
+    }
+
+    enum State {
+        OPEN,
+        CLOSED,
+        HALF_OPEN
+    }
+
+}