feat(circuit): support node circuit breaker (#2409)

* feat(circuit): add circuit object storage

Signed-off-by: Robin Han <[email protected]>

* fix(circuit): fix code review comment

Signed-off-by: Robin Han <[email protected]>

---------

Signed-off-by: Robin Han <[email protected]>

Author: superhx

Diff for: clients/src/main/java/org/apache/kafka/common/errors/s3/NodeLockedException.java

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2024, AutoMQ HK Limited.
+ *
+ * The use of this file is governed by the Business Source License,
+ * as detailed in the file "/LICENSE.S3Stream" included in this repository.
+ *
+ * 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.common.errors.s3;
+
+import org.apache.kafka.common.errors.ApiException;
+
+public class NodeLockedException extends ApiException {
+
+    public NodeLockedException(String message) {
+        super(message);
+    }
+
+}

Diff for: clients/src/main/java/org/apache/kafka/common/errors/s3/ObjectNotCommittedException.java

@@ -0,0 +1,20 @@
+/*
+ * Copyright 2024, AutoMQ HK Limited.
+ *
+ * The use of this file is governed by the Business Source License,
+ * as detailed in the file "/LICENSE.S3Stream" included in this repository.
+ *
+ * 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.common.errors.s3;
+
+import org.apache.kafka.common.errors.ApiException;
+
+public class ObjectNotCommittedException extends ApiException {
+    public ObjectNotCommittedException(String message) {
+        super(message);
+    }
+}

Diff for: clients/src/main/java/org/apache/kafka/common/protocol/Errors.java

@@ -151,6 +151,8 @@
 import org.apache.kafka.common.errors.s3.NodeEpochExpiredException;
 import org.apache.kafka.common.errors.s3.NodeEpochNotExistException;
 import org.apache.kafka.common.errors.s3.NodeFencedException;
+import org.apache.kafka.common.errors.s3.NodeLockedException;
+import org.apache.kafka.common.errors.s3.ObjectNotCommittedException;
 import org.apache.kafka.common.errors.s3.ObjectNotExistException;
 import org.apache.kafka.common.errors.s3.OffsetNotMatchedException;
 import org.apache.kafka.common.errors.s3.RedundantOperationException;
@@ -435,6 +437,8 @@ public enum Errors {
     KEY_EXIST(512, "The key already exists.", KeyExistException::new),
     KEY_NOT_EXIST(513, "The key does not exist.", ObjectNotExistException::new),
     NODE_FENCED(514, "The node is fenced.", NodeFencedException::new),
+    NODE_LOCKED(515, "The node is locked", NodeLockedException::new),
+    OBJECT_NOT_COMMITED(516, "The object is not commited.", ObjectNotCommittedException::new),
     STREAM_INNER_ERROR(599, "The stream inner error.", StreamInnerErrorException::new),
     // AutoMQ inject end
 

Diff for: core/src/main/java/kafka/autobalancer/AutoBalancerManager.java

@@ -75,6 +75,7 @@ protected void init() {
 
         this.anomalyDetector = new AnomalyDetectorImpl(config.originals(),
             new LogContext(String.format("[AnomalyDetector id=%d] ", nodeId)), clusterModel, actionExecutorService);
+        ((AnomalyDetectorImpl) this.anomalyDetector).lockedNodes(() -> quorumController.nodeControlManager().lockedNodes());
 
         this.reconfigurables.add(anomalyDetector);
 

Diff for: core/src/main/java/kafka/autobalancer/detector/AnomalyDetectorImpl.java

@@ -44,6 +44,7 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReentrantLock;
+import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 public class AnomalyDetectorImpl extends AbstractAnomalyDetector implements LeaderChangeListener {
@@ -56,6 +57,7 @@ public class AnomalyDetectorImpl extends AbstractAnomalyDetector implements Lead
     private final Lock configChangeLock = new ReentrantLock();
     private List<Goal> goalsByPriority;
     private Set<Integer> excludedBrokers;
+    private Supplier<Set<Integer>> lockedNodes = Collections::emptySet;
     private Set<String> excludedTopics;
     private long detectInterval;
     private long maxTolerateMetricsDelayMs;
@@ -266,6 +268,10 @@ public void detect() {
         this.executorService.schedule(this::detect, nextExecutionDelay, TimeUnit.MILLISECONDS);
     }
 
+    public void lockedNodes(Supplier<Set<Integer>> lockedNodes) {
+        this.lockedNodes = lockedNodes;
+    }
+
     private boolean isRunnable() {
         return this.running.get() && this.isLeader;
     }
@@ -283,6 +289,7 @@ long detect0() throws Exception {
         try {
             detectInterval = this.detectInterval;
             excludedBrokers = new HashSet<>(this.excludedBrokers);
+            excludedBrokers.addAll(lockedNodes.get());
             excludedTopics = new HashSet<>(this.excludedTopics);
             maxTolerateMetricsDelayMs = this.maxTolerateMetricsDelayMs;
             maxExecutionConcurrency = this.executionConcurrency;

Diff for: core/src/main/scala/kafka/controller/streamaspect/client/s3/StreamClientFactory.java

@@ -18,23 +18,27 @@
 import org.apache.kafka.controller.stream.StreamClient;
 
 import com.automq.stream.s3.Config;
+import com.automq.stream.s3.operator.ObjectStorage;
 import com.automq.stream.s3.operator.ObjectStorageFactory;
 
+import static com.automq.stream.s3.operator.ObjectStorageFactory.EXTENSION_TYPE_BACKGROUND;
+import static com.automq.stream.s3.operator.ObjectStorageFactory.EXTENSION_TYPE_KEY;
+
 public class StreamClientFactory {
 
     /**
      * This method will be called by {@link StreamClientFactoryProxy}
      */
     public static StreamClient get(Context context) {
         Config streamConfig = ConfigUtils.to(context.kafkaConfig);
+        ObjectStorage objectStorage = ObjectStorageFactory.instance().builder()
+            .buckets(streamConfig.dataBuckets())
+            .tagging(streamConfig.objectTagging())
+            .extension(EXTENSION_TYPE_KEY, EXTENSION_TYPE_BACKGROUND)
+            .build();
         return StreamClient.builder()
             .streamConfig(streamConfig)
-            .objectStorage(
-                ObjectStorageFactory.instance()
-                    .builder(streamConfig.dataBuckets().get(0))
-                    .tagging(streamConfig.objectTagging())
-                    .build()
-            )
+            .objectStorage(objectStorage)
             .build();
     }
 }

Diff for: core/src/main/scala/kafka/log/stream/s3/DefaultS3Client.java

@@ -65,6 +65,10 @@
 import java.util.concurrent.CompletableFuture;
 import java.util.concurrent.TimeUnit;
 
+import static com.automq.stream.s3.operator.ObjectStorageFactory.EXTENSION_TYPE_BACKGROUND;
+import static com.automq.stream.s3.operator.ObjectStorageFactory.EXTENSION_TYPE_KEY;
+import static com.automq.stream.s3.operator.ObjectStorageFactory.EXTENSION_TYPE_MAIN;
+
 public class DefaultS3Client implements Client {
     private static final Logger LOGGER = LoggerFactory.getLogger(DefaultS3Client.class);
     protected final Config config;
@@ -74,8 +78,8 @@ public class DefaultS3Client implements Client {
 
     protected ControllerRequestSender requestSender;
 
-    protected ObjectStorage objectStorage;
-    protected ObjectStorage compactionObjectStorage;
+    protected ObjectStorage mainObjectStorage;
+    protected ObjectStorage backgroundObjectStorage;
 
     protected WriteAheadLog writeAheadLog;
     protected StorageFailureHandlerChain storageFailureHandlerChain;
@@ -109,7 +113,6 @@ public DefaultS3Client(BrokerServer brokerServer, Config config) {
 
     @Override
     public void start() {
-        BucketURI dataBucket = config.dataBuckets().get(0);
         long refillToken = (long) (config.networkBaselineBandwidth() * ((double) config.refillPeriodMs() / 1000));
         if (refillToken <= 0) {
             throw new IllegalArgumentException(String.format("refillToken must be greater than 0, bandwidth: %d, refill period: %dms",
@@ -127,34 +130,30 @@ public void start() {
         S3StreamMetricsManager.registerNetworkAvailableBandwidthSupplier(AsyncNetworkBandwidthLimiter.Type.OUTBOUND, () ->
             config.networkBaselineBandwidth() - (long) networkOutboundRate.derive(
                 TimeUnit.NANOSECONDS.toSeconds(System.nanoTime()), NetworkStats.getInstance().networkOutboundUsageTotal().get()));
-        this.objectStorage = ObjectStorageFactory.instance().builder(dataBucket).tagging(config.objectTagging())
-            .inboundLimiter(networkInboundLimiter).outboundLimiter(networkOutboundLimiter).readWriteIsolate(true)
-            .threadPrefix("dataflow").build();
-        if (!objectStorage.readinessCheck()) {
+
+        this.localIndexCache = new LocalStreamRangeIndexCache();
+        this.objectReaderFactory = new DefaultObjectReaderFactory(() -> this.mainObjectStorage);
+        this.metadataManager = new StreamMetadataManager(brokerServer, config.nodeId(), objectReaderFactory, localIndexCache);
+        this.requestSender = new ControllerRequestSender(brokerServer, new ControllerRequestSender.RetryPolicyContext(config.controllerRequestRetryMaxCount(),
+            config.controllerRequestRetryBaseDelayMs()));
+        this.streamManager = newStreamManager(config.nodeId(), config.nodeEpoch(), false);
+        this.objectManager = newObjectManager(config.nodeId(), config.nodeEpoch(), false);
+        this.mainObjectStorage = newMainObjectStorage();
+        if (!mainObjectStorage.readinessCheck()) {
             throw new IllegalArgumentException(String.format("%s is not ready", config.dataBuckets()));
         }
-        this.compactionObjectStorage = ObjectStorageFactory.instance().builder(dataBucket).tagging(config.objectTagging())
-            .inboundLimiter(networkInboundLimiter).outboundLimiter(networkOutboundLimiter)
-            .threadPrefix("compaction").build();
-        ControllerRequestSender.RetryPolicyContext retryPolicyContext = new ControllerRequestSender.RetryPolicyContext(config.controllerRequestRetryMaxCount(),
-            config.controllerRequestRetryBaseDelayMs());
-        localIndexCache = new LocalStreamRangeIndexCache();
-        localIndexCache.init(config.nodeId(), objectStorage);
+        this.backgroundObjectStorage = newBackgroundObjectStorage();
+        localIndexCache.init(config.nodeId(), backgroundObjectStorage);
         localIndexCache.start();
-        this.objectReaderFactory = new DefaultObjectReaderFactory(objectStorage);
-        this.metadataManager = new StreamMetadataManager(brokerServer, config.nodeId(), objectReaderFactory, localIndexCache);
-        this.requestSender = new ControllerRequestSender(brokerServer, retryPolicyContext);
-        this.streamManager = newStreamManager(config.nodeId(), config.nodeEpoch(), false);
         this.streamManager.setStreamCloseHook(streamId -> localIndexCache.uploadOnStreamClose());
-        this.objectManager = newObjectManager(config.nodeId(), config.nodeEpoch(), false);
         this.objectManager.setCommitStreamSetObjectHook(localIndexCache::updateIndexFromRequest);
-        this.blockCache = new StreamReaders(this.config.blockCacheSize(), objectManager, objectStorage, objectReaderFactory);
-        this.compactionManager = new CompactionManager(this.config, this.objectManager, this.streamManager, compactionObjectStorage);
+        this.blockCache = new StreamReaders(this.config.blockCacheSize(), objectManager, mainObjectStorage, objectReaderFactory);
+        this.compactionManager = new CompactionManager(this.config, this.objectManager, this.streamManager, backgroundObjectStorage);
         this.writeAheadLog = buildWAL();
         this.storageFailureHandlerChain = new StorageFailureHandlerChain();
         this.storage = newS3Storage();
         // stream object compactions share the same object storage with stream set object compactions
-        this.streamClient = new S3StreamClient(this.streamManager, this.storage, this.objectManager, compactionObjectStorage, this.config, networkInboundLimiter, networkOutboundLimiter);
+        this.streamClient = new S3StreamClient(this.streamManager, this.storage, this.objectManager, backgroundObjectStorage, this.config, networkInboundLimiter, networkOutboundLimiter);
         storageFailureHandlerChain.addHandler(new ForceCloseStorageFailureHandler(streamClient));
         storageFailureHandlerChain.addHandler(new HaltStorageFailureHandler());
         this.streamClient.registerStreamLifeCycleListener(localIndexCache);
@@ -221,6 +220,22 @@ protected WriteAheadLog buildWAL() {
         }
     }
 
+    protected ObjectStorage newMainObjectStorage() {
+        return ObjectStorageFactory.instance().builder()
+            .buckets(config.dataBuckets())
+            .extension(EXTENSION_TYPE_KEY, EXTENSION_TYPE_MAIN)
+            .readWriteIsolate(true)
+            .build();
+    }
+
+    protected ObjectStorage newBackgroundObjectStorage() {
+        return ObjectStorageFactory.instance().builder()
+            .buckets(config.dataBuckets())
+            .extension(EXTENSION_TYPE_KEY, EXTENSION_TYPE_BACKGROUND)
+            .readWriteIsolate(false)
+            .build();
+    }
+
     protected StreamManager newStreamManager(int nodeId, long nodeEpoch, boolean failoverMode) {
         return new ControllerStreamManager(this.metadataManager, this.requestSender, nodeId, nodeEpoch,
             this::getAutoMQVersion, failoverMode);
@@ -232,7 +247,7 @@ protected ObjectManager newObjectManager(int nodeId, long nodeEpoch, boolean fai
     }
 
     protected S3Storage newS3Storage() {
-        return new S3Storage(config, writeAheadLog, streamManager, objectManager, blockCache, objectStorage, storageFailureHandlerChain);
+        return new S3Storage(config, writeAheadLog, streamManager, objectManager, blockCache, mainObjectStorage, storageFailureHandlerChain);
     }
 
     protected Failover failover() {

Diff for: core/src/main/scala/kafka/log/stream/s3/objects/ControllerObjectManager.java

@@ -34,6 +34,8 @@
 
 import com.automq.stream.s3.compact.CompactOperations;
 import com.automq.stream.s3.exceptions.AutoMQException;
+import com.automq.stream.s3.exceptions.CompactedObjectsNotFoundException;
+import com.automq.stream.s3.exceptions.ObjectNotCommittedException;
 import com.automq.stream.s3.metadata.S3ObjectMetadata;
 import com.automq.stream.s3.objects.CommitStreamSetObjectHook;
 import com.automq.stream.s3.objects.CommitStreamSetObjectRequest;
@@ -180,7 +182,7 @@ public Builder toRequestBuilder() {
                     throw Errors.forCode(resp.errorCode()).exception();
                 case OBJECT_NOT_EXIST:
                 case COMPACTED_OBJECTS_NOT_FOUND:
-                    throw code.exception();
+                    throw new CompactedObjectsNotFoundException();
                 default:
                     LOGGER.error("Error while committing stream set object: {}, code: {}, retry later", request, code);
                     return ResponseHandleResult.withRetry();
@@ -231,7 +233,9 @@ public Builder toRequestBuilder() {
                     throw Errors.forCode(resp.errorCode()).exception();
                 case OBJECT_NOT_EXIST:
                 case COMPACTED_OBJECTS_NOT_FOUND:
-                    throw code.exception();
+                    throw new CompactedObjectsNotFoundException();
+                case OBJECT_NOT_COMMITED:
+                    throw new ObjectNotCommittedException();
                 case STREAM_NOT_EXIST:
                 case STREAM_FENCED:
                     LOGGER.warn("Stream fenced or not exist: {}, code: {}", request, Errors.forCode(resp.errorCode()));

Diff for: core/src/main/scala/kafka/log/stream/s3/streams/ControllerStreamManager.java

@@ -259,6 +259,9 @@ public String toString() {
                 case STREAM_NOT_CLOSED:
                     logger.warn("open stream fail: {}, epoch {}, code: STREAM_NOT_CLOSED, retry later", streamId, epoch);
                     return ResponseHandleResult.withRetry();
+                case NODE_LOCKED:
+                    logger.warn("[NODE_LOCKED]open stream fail: {}, epoch {}", streamId, epoch);
+                    throw code.exception();
                 default:
                     logger.error("Error while opening stream: {}, epoch {}, code: {}, retry later", streamId, epoch, code);
                     return ResponseHandleResult.withRetry();

Diff for: core/src/main/scala/kafka/log/streamaspect/DefaultOpenStreamChecker.java

@@ -22,19 +22,23 @@
 
 import com.automq.stream.s3.metadata.StreamState;
 
+import java.util.Arrays;
+
 public class DefaultOpenStreamChecker implements OpenStreamChecker {
+    private final int nodeId;
     private final KRaftMetadataCache metadataCache;
 
-    public DefaultOpenStreamChecker(KRaftMetadataCache metadataCache) {
+    public DefaultOpenStreamChecker(int nodeId, KRaftMetadataCache metadataCache) {
+        this.nodeId = nodeId;
         this.metadataCache = metadataCache;
     }
 
     @Override
     public boolean check(Uuid topicId, int partition, long streamId, long epoch) throws StreamFencedException {
-        return metadataCache.safeRun(image -> DefaultOpenStreamChecker.check(image, topicId, partition, streamId, epoch));
+        return metadataCache.safeRun(image -> DefaultOpenStreamChecker.check(image, topicId, partition, streamId, epoch, nodeId));
     }
 
-    public static boolean check(MetadataImage image, Uuid topicId, int partition, long streamId, long epoch) throws StreamFencedException {
+    public static boolean check(MetadataImage image, Uuid topicId, int partition, long streamId, long epoch, int currentNodeId) throws StreamFencedException {
         // When ABA reassign happens:
         // 1. Assign P0 to broker0 with epoch=0, broker0 opens the partition
         // 2. Assign P0 to broker1 with epoch=1, broker1 waits for the partition to be closed
@@ -52,6 +56,10 @@ public static boolean check(MetadataImage image, Uuid topicId, int partition, lo
         if (currentEpoch > epoch) {
             throw new StreamFencedException(String.format("partition=%s-%d with epoch=%d is fenced by new leader epoch=%d", topicId, partition, epoch, currentEpoch));
         }
+        if (!contains(partitionImage.isr, currentNodeId)) {
+            throw new StreamFencedException(String.format("partition=%s-%d with epoch=%d move to other nodes %s", topicId, partition, epoch, Arrays.toString(partitionImage.isr)));
+        }
+
         S3StreamMetadataImage stream = image.streamsMetadata().getStreamMetadata(streamId);
         if (stream == null) {
             throw new StreamFencedException(String.format("streamId=%d cannot be found, it may be deleted or not created yet", streamId));
@@ -60,4 +68,16 @@ public static boolean check(MetadataImage image, Uuid topicId, int partition, lo
             throw new StreamFencedException(String.format("streamId=%d with epoch=%d is fenced by new leader epoch=%d", streamId, epoch, stream.getEpoch()));
         return StreamState.CLOSED.equals(stream.state());
     }
+
+    private static boolean contains(int[] isr, int nodeId) {
+        if (isr == null) {
+            return false;
+        }
+        for (int replica : isr) {
+            if (replica == nodeId) {
+                return true;
+            }
+        }
+        return false;
+    }
 }

Diff for: core/src/main/scala/kafka/log/streamaspect/ElasticLogManager.scala

@@ -133,7 +133,7 @@ object ElasticLogManager {
     context.config = config
     context.brokerServer = broker
     val openStreamChecker = if (broker != null) {
-      new DefaultOpenStreamChecker(broker.metadataCache)
+      new DefaultOpenStreamChecker(config.nodeId, broker.metadataCache)
     } else {
       OpenStreamChecker.NOOP
     }

Diff for: metadata/src/main/java/org/apache/kafka/controller/QuorumController.java

@@ -145,7 +145,7 @@
 import org.apache.kafka.controller.errors.ControllerExceptions;
 import org.apache.kafka.controller.errors.EventHandlerExceptionInfo;
 import org.apache.kafka.controller.metrics.QuorumControllerMetrics;
-import org.apache.kafka.controller.stream.DefaultNodeRuntimeInfoGetter;
+import org.apache.kafka.controller.stream.DefaultNodeRuntimeInfoManager;
 import org.apache.kafka.controller.stream.KVControlManager;
 import org.apache.kafka.controller.stream.NodeControlManager;
 import org.apache.kafka.controller.stream.S3ObjectControlManager;
@@ -2165,7 +2165,7 @@ private QuorumController(
         this.streamControlManager = new StreamControlManager(this, snapshotRegistry, logContext,
                 this.s3ObjectControlManager, clusterControl, featureControl, replicationControl);
         this.topicDeletionManager = new TopicDeletionManager(snapshotRegistry, this, streamControlManager, kvControlManager);
-        this.nodeControlManager = new NodeControlManager(snapshotRegistry, new DefaultNodeRuntimeInfoGetter(clusterControl, streamControlManager));
+        this.nodeControlManager = new NodeControlManager(snapshotRegistry, new DefaultNodeRuntimeInfoManager(clusterControl, streamControlManager));
         this.extension = extension.apply(this);
 
         // set the nodeControlManager here to avoid circular dependency