diff --git a/core/src/main/java/kafka/automq/AutoMQConfig.java b/core/src/main/java/kafka/automq/AutoMQConfig.java index 7082e27677..2b01ac1fad 100644 --- a/core/src/main/java/kafka/automq/AutoMQConfig.java +++ b/core/src/main/java/kafka/automq/AutoMQConfig.java @@ -43,6 +43,7 @@ import static org.apache.kafka.common.config.ConfigDef.Importance.HIGH; import static org.apache.kafka.common.config.ConfigDef.Importance.LOW; import static org.apache.kafka.common.config.ConfigDef.Importance.MEDIUM; +import static org.apache.kafka.common.config.ConfigDef.Range.atLeast; import static org.apache.kafka.common.config.ConfigDef.Range.between; import static org.apache.kafka.common.config.ConfigDef.Type.BOOLEAN; import static org.apache.kafka.common.config.ConfigDef.Type.INT; @@ -202,6 +203,12 @@ public class AutoMQConfig { public static final String RETRY_STORM_BACKOFF_ENABLED_DOC = "Whether retry storm delayed response backoff is enabled"; public static final boolean RETRY_STORM_BACKOFF_ENABLED_DEFAULT = false; + public static final String KAFKA_GO_METADATA_COMPATIBILITY_ENABLED_CONFIG = + "automq.kafka-go.metadata.compatibility.enabled"; + public static final String KAFKA_GO_METADATA_COMPATIBILITY_ENABLED_DOC = + "Whether kafka-go Metadata compatibility mode is enabled"; + public static final boolean KAFKA_GO_METADATA_COMPATIBILITY_ENABLED_DEFAULT = true; + public static final String RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG = "automq.retry.storm.backoff.max.delay.ms"; public static final String RETRY_STORM_BACKOFF_MAX_DELAY_MS_DOC = "The maximum retry storm delayed response time in milliseconds, from 0 to 10000"; public static final long RETRY_STORM_BACKOFF_MAX_DELAY_MS_DEFAULT = 1000L; @@ -310,9 +317,10 @@ public static void define(ConfigDef configDef) { .define(AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_CONFIG, PASSWORD, null, HIGH, AutoMQConfig.S3_TELEMETRY_METRICS_EXPORTER_URI_DOC) .define(AutoMQConfig.S3_TELEMETRY_METRICS_BASE_LABELS_CONFIG, STRING, null, MEDIUM, AutoMQConfig.S3_TELEMETRY_METRICS_BASE_LABELS_DOC) .define(AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG, BOOLEAN, AutoMQConfig.S3_BACK_PRESSURE_ENABLED_DEFAULT, MEDIUM, AutoMQConfig.S3_BACK_PRESSURE_ENABLED_DOC) - .define(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG, LONG, AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_DEFAULT, MEDIUM, AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_DOC) + .define(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG, LONG, AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_DEFAULT, atLeast(0), MEDIUM, AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_DOC) .define(AutoMQConfig.RETRY_STORM_BACKOFF_ENABLED_CONFIG, BOOLEAN, AutoMQConfig.RETRY_STORM_BACKOFF_ENABLED_DEFAULT, MEDIUM, AutoMQConfig.RETRY_STORM_BACKOFF_ENABLED_DOC) .define(AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG, LONG, AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_DEFAULT, between(0, AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_MAX), MEDIUM, AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_DOC) + .define(AutoMQConfig.KAFKA_GO_METADATA_COMPATIBILITY_ENABLED_CONFIG, BOOLEAN, AutoMQConfig.KAFKA_GO_METADATA_COMPATIBILITY_ENABLED_DEFAULT, MEDIUM, AutoMQConfig.KAFKA_GO_METADATA_COMPATIBILITY_ENABLED_DOC) .define(AutoMQConfig.ZONE_ROUTER_CHANNELS_CONFIG, ConfigDef.Type.STRING, null, ConfigDef.Importance.HIGH, AutoMQConfig.ZONE_ROUTER_CHANNELS_DOC) // Deprecated config start .define(AutoMQConfig.S3_ENDPOINT_CONFIG, STRING, null, HIGH, AutoMQConfig.S3_ENDPOINT_DOC) diff --git a/core/src/main/java/kafka/automq/backpressure/BackPressureConfig.java b/core/src/main/java/kafka/automq/backpressure/BackPressureConfig.java index d8c7c946ce..0404e7f18a 100644 --- a/core/src/main/java/kafka/automq/backpressure/BackPressureConfig.java +++ b/core/src/main/java/kafka/automq/backpressure/BackPressureConfig.java @@ -59,16 +59,6 @@ public BackPressureConfig(boolean enabled, long cooldownMs) { this.cooldownMs = cooldownMs; } - public static void validate(Map raw) throws ConfigException { - Map configs = new HashMap<>(raw); - if (configs.containsKey(AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG)) { - ConfigUtils.getBoolean(configs, AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG); - } - if (configs.containsKey(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG)) { - validateCooldownMs(ConfigUtils.getLong(configs, AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG)); - } - } - public static void validateCooldownMs(long cooldownMs) throws ConfigException { if (cooldownMs < 0) { throw new ConfigException(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG, cooldownMs, "The cooldown time must be non-negative."); @@ -76,12 +66,11 @@ public static void validateCooldownMs(long cooldownMs) throws ConfigException { } public void update(Map raw) { - Map configs = new HashMap<>(raw); - if (configs.containsKey(AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG)) { - this.enabled = ConfigUtils.getBoolean(configs, AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG); + if (raw.containsKey(AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG)) { + this.enabled = (Boolean) raw.get(AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG); } - if (configs.containsKey(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG)) { - this.cooldownMs = ConfigUtils.getLong(configs, AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG); + if (raw.containsKey(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG)) { + this.cooldownMs = (Long) raw.get(AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG); } } diff --git a/core/src/main/java/kafka/automq/backpressure/DefaultBackPressureManager.java b/core/src/main/java/kafka/automq/backpressure/DefaultBackPressureManager.java index 7f129e49ae..a5995095e5 100644 --- a/core/src/main/java/kafka/automq/backpressure/DefaultBackPressureManager.java +++ b/core/src/main/java/kafka/automq/backpressure/DefaultBackPressureManager.java @@ -181,7 +181,6 @@ public Set reconfigurableConfigs() { @Override public void validateReconfiguration(Map configs) throws ConfigException { - BackPressureConfig.validate(configs); } @Override diff --git a/core/src/main/java/kafka/automq/metadata/KafkaGoMetadataResponseRewriter.java b/core/src/main/java/kafka/automq/metadata/KafkaGoMetadataResponseRewriter.java new file mode 100644 index 0000000000..6f593d2307 --- /dev/null +++ b/core/src/main/java/kafka/automq/metadata/KafkaGoMetadataResponseRewriter.java @@ -0,0 +1,89 @@ +/* + * Copyright 2026, AutoMQ HK Limited. + * + * 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 kafka.automq.metadata; + +import kafka.automq.AutoMQConfig; +import kafka.server.KafkaConfig; + +import org.apache.kafka.common.Reconfigurable; +import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseTopic; +import org.apache.kafka.common.protocol.Errors; + +import java.util.Collection; +import java.util.Map; +import java.util.Set; + +/** + * Rewrites unavailable kafka-go partition metadata to route requests to the controller-selected first replica. + */ +public final class KafkaGoMetadataResponseRewriter implements Reconfigurable { + + public static final Set RECONFIGURABLE_CONFIGS = + Set.of(AutoMQConfig.KAFKA_GO_METADATA_COMPATIBILITY_ENABLED_CONFIG); + + private static final String KAFKA_GO_CLIENT_ID = "github.com/segmentio/kafka-go"; + private volatile boolean enabled; + + public KafkaGoMetadataResponseRewriter(KafkaConfig config) { + this.enabled = config.kafkaGoMetadataCompatibilityEnabled(); + } + + @Override + public Set reconfigurableConfigs() { + return RECONFIGURABLE_CONFIGS; + } + + @Override + public void validateReconfiguration(Map configs) { + } + + @Override + public void reconfigure(Map configs) { + if (configs.containsKey(AutoMQConfig.KAFKA_GO_METADATA_COMPATIBILITY_ENABLED_CONFIG)) { + enabled = (Boolean) configs.get(AutoMQConfig.KAFKA_GO_METADATA_COMPATIBILITY_ENABLED_CONFIG); + } + } + + @Override + public void configure(Map configs) { + } + + /** + * Rewrites each leader-unavailable partition with replicas for clients using kafka-go's default client ID. + * All partition fields other than the error code and leader ID are preserved. + * + * @param clientId request header client ID + * @param topics mutable Metadata response topics + */ + public void rewrite(String clientId, Collection topics) { + if (!enabled + || clientId == null + || (!clientId.isEmpty() && !clientId.contains(KAFKA_GO_CLIENT_ID))) { + return; + } + topics.forEach(topic -> topic.partitions().forEach(partition -> { + if (partition.errorCode() == Errors.LEADER_NOT_AVAILABLE.code() + && !partition.replicaNodes().isEmpty()) { + partition.setErrorCode(Errors.NONE.code()); + partition.setLeaderId(partition.replicaNodes().get(0)); + } + })); + } +} diff --git a/core/src/main/java/kafka/automq/retrystorm/RetryStormBackoffConfig.java b/core/src/main/java/kafka/automq/retrystorm/RetryStormBackoffConfig.java index 9fe6e0cd11..ce6f2ee193 100644 --- a/core/src/main/java/kafka/automq/retrystorm/RetryStormBackoffConfig.java +++ b/core/src/main/java/kafka/automq/retrystorm/RetryStormBackoffConfig.java @@ -76,19 +76,6 @@ public static RetryStormBackoffConfig from(Map raw) { ); } - /** - * Validates retry storm keys present in a dynamic broker config update. - */ - public static void validate(Map raw) throws ConfigException { - Map configs = new HashMap<>(raw); - if (configs.containsKey(AutoMQConfig.RETRY_STORM_BACKOFF_ENABLED_CONFIG)) { - ConfigUtils.getBoolean(configs, AutoMQConfig.RETRY_STORM_BACKOFF_ENABLED_CONFIG); - } - if (configs.containsKey(AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG)) { - validateMaxDelayMs(ConfigUtils.getLong(configs, AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG)); - } - } - /** * Rejects max delay values outside the bounded delayed-response range. */ @@ -104,14 +91,11 @@ public static void validateMaxDelayMs(long maxDelayMs) throws ConfigException { * Applies a partial dynamic update; keys absent from {@code raw} keep their current runtime values. */ public void update(Map raw) { - Map configs = new HashMap<>(raw); - if (configs.containsKey(AutoMQConfig.RETRY_STORM_BACKOFF_ENABLED_CONFIG)) { - this.enabled = ConfigUtils.getBoolean(configs, AutoMQConfig.RETRY_STORM_BACKOFF_ENABLED_CONFIG); + if (raw.containsKey(AutoMQConfig.RETRY_STORM_BACKOFF_ENABLED_CONFIG)) { + this.enabled = (Boolean) raw.get(AutoMQConfig.RETRY_STORM_BACKOFF_ENABLED_CONFIG); } - if (configs.containsKey(AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG)) { - long nextMaxDelayMs = ConfigUtils.getLong(configs, AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG); - validateMaxDelayMs(nextMaxDelayMs); - this.maxDelayMs = nextMaxDelayMs; + if (raw.containsKey(AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG)) { + this.maxDelayMs = (Long) raw.get(AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG); } } diff --git a/core/src/main/java/kafka/automq/retrystorm/RetryStormBackoffManager.java b/core/src/main/java/kafka/automq/retrystorm/RetryStormBackoffManager.java index 30a94f003e..9a69ee1833 100644 --- a/core/src/main/java/kafka/automq/retrystorm/RetryStormBackoffManager.java +++ b/core/src/main/java/kafka/automq/retrystorm/RetryStormBackoffManager.java @@ -106,7 +106,6 @@ public Set reconfigurableConfigs() { */ @Override public void validateReconfiguration(Map configs) throws ConfigException { - RetryStormBackoffConfig.validate(configs); } /** diff --git a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala index 507e724fca..783b700791 100755 --- a/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala +++ b/core/src/main/scala/kafka/server/DynamicBrokerConfig.scala @@ -19,6 +19,7 @@ package kafka.server import kafka.autobalancer.config.{AutoBalancerControllerConfig, AutoBalancerMetricsReporterConfig} import kafka.automq.backpressure.BackPressureConfig +import kafka.automq.metadata.KafkaGoMetadataResponseRewriter import kafka.automq.retrystorm.RetryStormBackoffConfig import java.util @@ -105,7 +106,8 @@ object DynamicBrokerConfig { AutoBalancerControllerConfig.RECONFIGURABLE_CONFIGS.asScala ++ AutoBalancerMetricsReporterConfig.RECONFIGURABLE_CONFIGS.asScala ++ BackPressureConfig.RECONFIGURABLE_CONFIGS.asScala ++ - RetryStormBackoffConfig.RECONFIGURABLE_CONFIGS.asScala + RetryStormBackoffConfig.RECONFIGURABLE_CONFIGS.asScala ++ + KafkaGoMetadataResponseRewriter.RECONFIGURABLE_CONFIGS.asScala private val ClusterLevelListenerConfigs = Set(SocketServerConfigs.MAX_CONNECTIONS_CONFIG, SocketServerConfigs.MAX_CONNECTION_CREATION_RATE_CONFIG, SocketServerConfigs.NUM_NETWORK_THREADS_CONFIG) private val PerBrokerConfigs = (DynamicSecurityConfigs ++ DynamicListenerConfig.ReconfigurableConfigs).diff( diff --git a/core/src/main/scala/kafka/server/KafkaApis.scala b/core/src/main/scala/kafka/server/KafkaApis.scala index 32e86019d0..78bc066111 100644 --- a/core/src/main/scala/kafka/server/KafkaApis.scala +++ b/core/src/main/scala/kafka/server/KafkaApis.scala @@ -19,6 +19,7 @@ package kafka.server import kafka.api.ElectLeadersRequestOps import kafka.automq.interceptor.ClientIdMetadata +import kafka.automq.metadata.KafkaGoMetadataResponseRewriter import kafka.controller.ReplicaAssignment import kafka.coordinator.transaction.{InitProducerIdResult, TransactionCoordinator} import kafka.network.RequestChannel @@ -133,6 +134,8 @@ class KafkaApis(val requestChannel: RequestChannel, private val listOffsetSlowExecutor = KafkaApis.newListOffsetExecutor("kafka-apis-list-offset-slow-%d") private val listOffsetRequestExecutor = KafkaApis.newListOffsetExecutor("kafka-apis-list-offset-request-%d") private val listOffsetInflightPartitions = new AtomicInteger(0) + private val kafkaGoMetadataResponseRewriter = new KafkaGoMetadataResponseRewriter(config) + config.addReconfigurable(kafkaGoMetadataResponseRewriter) // AutoMQ inject end @@ -141,6 +144,7 @@ class KafkaApis(val requestChannel: RequestChannel, ThreadUtils.shutdownExecutor(listOffsetFastExecutor, 5, TimeUnit.SECONDS, logger.underlying) ThreadUtils.shutdownExecutor(listOffsetSlowExecutor, 5, TimeUnit.SECONDS, logger.underlying) ThreadUtils.shutdownExecutor(listOffsetRequestExecutor, 5, TimeUnit.SECONDS, logger.underlying) + config.removeReconfigurable(kafkaGoMetadataResponseRewriter) // AutoMQ inject end aclApis.close() info("Shutdown complete.") @@ -1500,6 +1504,10 @@ class KafkaApis(val requestChannel: RequestChannel, val completeTopicMetadata = unknownTopicIdsTopicMetadata ++ topicMetadata ++ unauthorizedForCreateTopicMetadata ++ unauthorizedForDescribeTopicMetadata + // AutoMQ inject start + kafkaGoMetadataResponseRewriter.rewrite(request.header.clientId, completeTopicMetadata.asJava) + // AutoMQ inject end + val brokers = metadataCache.getAliveBrokerNodes(request.context.listenerName) trace("Sending topic metadata %s and brokers %s for correlation id %d to client %s".format(completeTopicMetadata.mkString(","), diff --git a/core/src/main/scala/kafka/server/KafkaConfig.scala b/core/src/main/scala/kafka/server/KafkaConfig.scala index d35026b9e1..1c2ffe527e 100755 --- a/core/src/main/scala/kafka/server/KafkaConfig.scala +++ b/core/src/main/scala/kafka/server/KafkaConfig.scala @@ -802,6 +802,7 @@ class KafkaConfig private(doLog: Boolean, val props: util.Map[_, _]) val tableTopicSchemaRegistryUrl = getString(AutoMQConfig.TABLE_TOPIC_SCHEMA_REGISTRY_URL_CONFIG) val retryStormBackoffEnabled = getBoolean(AutoMQConfig.RETRY_STORM_BACKOFF_ENABLED_CONFIG) val retryStormBackoffMaxDelayMs = getLong(AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG) + def kafkaGoMetadataCompatibilityEnabled: Boolean = getBoolean(AutoMQConfig.KAFKA_GO_METADATA_COMPATIBILITY_ENABLED_CONFIG) // AutoMQ inject end /** Internal Configurations **/ diff --git a/core/src/test/java/kafka/automq/backpressure/DefaultBackPressureManagerTest.java b/core/src/test/java/kafka/automq/backpressure/DefaultBackPressureManagerTest.java index f2f426aa84..7d3334cb34 100644 --- a/core/src/test/java/kafka/automq/backpressure/DefaultBackPressureManagerTest.java +++ b/core/src/test/java/kafka/automq/backpressure/DefaultBackPressureManagerTest.java @@ -21,6 +21,9 @@ import kafka.automq.AutoMQConfig; +import org.apache.kafka.common.config.ConfigDef; +import org.apache.kafka.common.config.ConfigException; + import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; @@ -29,6 +32,7 @@ import java.util.concurrent.TimeUnit; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.ArgumentMatchers.anyLong; import static org.mockito.Mockito.doAnswer; @@ -50,6 +54,15 @@ public class DefaultBackPressureManagerTest { int schedulerScheduleCalled = 0; long schedulerScheduleDelay = 0; + @Test + public void testConfigDefRejectsNegativeCooldown() { + ConfigDef configDef = new ConfigDef(); + AutoMQConfig.define(configDef); + + assertThrows(ConfigException.class, () -> configDef.parse(Map.of( + AutoMQConfig.S3_BACK_PRESSURE_COOLDOWN_MS_CONFIG, -1L))); + } + @BeforeEach public void setup() { regulator = mock(Regulator.class); @@ -89,14 +102,14 @@ public void testDynamicConfig() { assertRegulatorCalled(0, 0); manager.reconfigure(Map.of( - AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG, "true" + AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG, true )); callChecker(sourceC, LoadLevel.NORMAL); callChecker(sourceB, LoadLevel.NORMAL); assertRegulatorCalled(1, 1); manager.reconfigure(Map.of( - AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG, "false" + AutoMQConfig.S3_BACK_PRESSURE_ENABLED_CONFIG, false )); callChecker(sourceC, LoadLevel.NORMAL); callChecker(sourceB, LoadLevel.HIGH); diff --git a/core/src/test/java/kafka/automq/metadata/KafkaGoMetadataResponseRewriterTest.java b/core/src/test/java/kafka/automq/metadata/KafkaGoMetadataResponseRewriterTest.java new file mode 100644 index 0000000000..347bd9ed19 --- /dev/null +++ b/core/src/test/java/kafka/automq/metadata/KafkaGoMetadataResponseRewriterTest.java @@ -0,0 +1,201 @@ +/* + * Copyright 2026, AutoMQ HK Limited. + * + * 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 kafka.automq.metadata; + +import kafka.automq.AutoMQConfig; +import kafka.server.KafkaConfig; + +import org.apache.kafka.common.message.MetadataResponseData.MetadataResponsePartition; +import org.apache.kafka.common.message.MetadataResponseData.MetadataResponseTopic; +import org.apache.kafka.common.protocol.Errors; + +import org.junit.jupiter.api.Tag; +import org.junit.jupiter.api.Test; + +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.when; + +/** + * Covers kafka-go-specific predictive leader rewrites in Metadata responses. + */ +@Tag("S3Unit") +public class KafkaGoMetadataResponseRewriterTest { + + private static final String KAFKA_GO_CLIENT_ID = "service/github.com/segmentio/kafka-go"; + + /** + * Given a kafka-go client and leader-unavailable partition with replicas, the first replica becomes the leader. + */ + @Test + public void testRewriteLeaderNotAvailableForKafkaGo() { + MetadataResponsePartition partition = partition(Errors.LEADER_NOT_AVAILABLE, -1, List.of(2, 1)); + + enabledRewriter().rewrite(KAFKA_GO_CLIENT_ID, List.of(topic(partition))); + + assertEquals(Errors.NONE.code(), partition.errorCode()); + assertEquals(2, partition.leaderId()); + } + + /** + * Given an empty client ID used by the default kafka-go Writer, leader-unavailable metadata is rewritten. + */ + @Test + public void testRewriteLeaderNotAvailableForEmptyClientId() { + MetadataResponsePartition partition = partition(Errors.LEADER_NOT_AVAILABLE, -1, List.of(2, 1)); + + enabledRewriter().rewrite("", List.of(topic(partition))); + + assertEquals(Errors.NONE.code(), partition.errorCode()); + assertEquals(2, partition.leaderId()); + } + + /** + * Given a non-kafka-go client, leader-unavailable metadata remains unchanged. + */ + @Test + public void testDoesNotRewriteOtherClients() { + MetadataResponsePartition partition = partition(Errors.LEADER_NOT_AVAILABLE, -1, List.of(2)); + + enabledRewriter().rewrite("apache-kafka-java", List.of(topic(partition))); + + assertEquals(Errors.LEADER_NOT_AVAILABLE.code(), partition.errorCode()); + assertEquals(-1, partition.leaderId()); + } + + /** + * Given a kafka-go client with no replicas, leader-unavailable metadata remains unchanged. + */ + @Test + public void testDoesNotRewriteEmptyReplicas() { + MetadataResponsePartition partition = partition(Errors.LEADER_NOT_AVAILABLE, -1, List.of()); + + enabledRewriter().rewrite(KAFKA_GO_CLIENT_ID, List.of(topic(partition))); + + assertEquals(Errors.LEADER_NOT_AVAILABLE.code(), partition.errorCode()); + assertEquals(-1, partition.leaderId()); + } + + /** + * Given a kafka-go client and a different partition error, metadata remains unchanged. + */ + @Test + public void testDoesNotRewriteOtherErrors() { + MetadataResponsePartition partition = partition(Errors.NOT_LEADER_OR_FOLLOWER, -1, List.of(2)); + + enabledRewriter().rewrite(KAFKA_GO_CLIENT_ID, List.of(topic(partition))); + + assertEquals(Errors.NOT_LEADER_OR_FOLLOWER.code(), partition.errorCode()); + assertEquals(-1, partition.leaderId()); + } + + /** + * Given mixed partition states, only eligible leader-unavailable partitions are rewritten. + */ + @Test + public void testRewritesOnlyEligiblePartitions() { + MetadataResponsePartition eligible = partition(Errors.LEADER_NOT_AVAILABLE, -1, List.of(2)); + MetadataResponsePartition emptyReplicas = partition(Errors.LEADER_NOT_AVAILABLE, -1, List.of()); + MetadataResponsePartition otherError = partition(Errors.UNKNOWN_TOPIC_OR_PARTITION, -1, List.of(1)); + + enabledRewriter().rewrite(KAFKA_GO_CLIENT_ID, + List.of(topic(eligible, emptyReplicas), topic(otherError))); + + assertEquals(Errors.NONE.code(), eligible.errorCode()); + assertEquals(2, eligible.leaderId()); + assertEquals(Errors.LEADER_NOT_AVAILABLE.code(), emptyReplicas.errorCode()); + assertEquals(Errors.UNKNOWN_TOPIC_OR_PARTITION.code(), otherError.errorCode()); + } + + /** + * Given an eligible rewrite, leader epoch and replica health fields are preserved. + */ + @Test + public void testPreservesOtherPartitionFields() { + MetadataResponsePartition partition = partition(Errors.LEADER_NOT_AVAILABLE, -1, List.of(2, 1)) + .setPartitionIndex(3) + .setLeaderEpoch(9) + .setIsrNodes(List.of(1)) + .setOfflineReplicas(List.of(2)); + + enabledRewriter().rewrite(KAFKA_GO_CLIENT_ID, List.of(topic(partition))); + + assertEquals(3, partition.partitionIndex()); + assertEquals(9, partition.leaderEpoch()); + assertEquals(List.of(2, 1), partition.replicaNodes()); + assertEquals(List.of(1), partition.isrNodes()); + assertEquals(List.of(2), partition.offlineReplicas()); + } + + /** + * Given compatibility mode is disabled, kafka-go metadata remains unchanged. + */ + @Test + public void testDoesNotRewriteWhenCompatibilityModeDisabled() { + MetadataResponsePartition partition = partition(Errors.LEADER_NOT_AVAILABLE, -1, List.of(2, 1)); + + KafkaConfig config = mock(KafkaConfig.class); + when(config.kafkaGoMetadataCompatibilityEnabled()).thenReturn(false); + + new KafkaGoMetadataResponseRewriter(config).rewrite(KAFKA_GO_CLIENT_ID, List.of(topic(partition))); + + assertEquals(Errors.LEADER_NOT_AVAILABLE.code(), partition.errorCode()); + assertEquals(-1, partition.leaderId()); + } + + /** + * Given compatibility mode is disabled dynamically, the existing rewriter observes the new value. + */ + @Test + public void testReadsCompatibilityModeForEachRewrite() { + KafkaConfig config = mock(KafkaConfig.class); + when(config.kafkaGoMetadataCompatibilityEnabled()).thenReturn(true); + KafkaGoMetadataResponseRewriter rewriter = new KafkaGoMetadataResponseRewriter(config); + MetadataResponsePartition first = partition(Errors.LEADER_NOT_AVAILABLE, -1, List.of(2)); + MetadataResponsePartition second = partition(Errors.LEADER_NOT_AVAILABLE, -1, List.of(2)); + + rewriter.rewrite(KAFKA_GO_CLIENT_ID, List.of(topic(first))); + rewriter.reconfigure(Map.of(AutoMQConfig.KAFKA_GO_METADATA_COMPATIBILITY_ENABLED_CONFIG, false)); + rewriter.rewrite(KAFKA_GO_CLIENT_ID, List.of(topic(second))); + + assertEquals(Errors.NONE.code(), first.errorCode()); + assertEquals(Errors.LEADER_NOT_AVAILABLE.code(), second.errorCode()); + } + + private static KafkaGoMetadataResponseRewriter enabledRewriter() { + KafkaConfig config = mock(KafkaConfig.class); + when(config.kafkaGoMetadataCompatibilityEnabled()).thenReturn(true); + return new KafkaGoMetadataResponseRewriter(config); + } + + private static MetadataResponsePartition partition(Errors error, int leaderId, List replicas) { + return new MetadataResponsePartition() + .setErrorCode(error.code()) + .setLeaderId(leaderId) + .setReplicaNodes(replicas); + } + + private static MetadataResponseTopic topic(MetadataResponsePartition... partitions) { + return new MetadataResponseTopic().setPartitions(List.of(partitions)); + } +} diff --git a/core/src/test/java/kafka/automq/retrystorm/RetryStormBackoffConfigTest.java b/core/src/test/java/kafka/automq/retrystorm/RetryStormBackoffConfigTest.java index fa9463d021..c4f327c43d 100644 --- a/core/src/test/java/kafka/automq/retrystorm/RetryStormBackoffConfigTest.java +++ b/core/src/test/java/kafka/automq/retrystorm/RetryStormBackoffConfigTest.java @@ -53,16 +53,6 @@ public void testConfigDefDefaults() { assertEquals(1000L, defaults.get(AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG)); } - /** - * Given a dynamic config update, validation rejects a negative max delay before runtime update. - */ - @Test - public void testValidateRejectsNegativeMaxDelayMs() { - assertThrows(ConfigException.class, () -> RetryStormBackoffConfig.validate(Map.of( - AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG, -1L - ))); - } - /** * Given static broker config parsing, ConfigDef rejects negative max delay values. */ @@ -89,16 +79,6 @@ public void testConfigDefRejectsMaxDelayAboveUpperBound() { ))); } - /** - * Given a dynamic config update, validation rejects values above the retry storm delay upper bound. - */ - @Test - public void testValidateRejectsMaxDelayAboveUpperBound() { - assertThrows(ConfigException.class, () -> RetryStormBackoffConfig.validate(Map.of( - AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG, 10001L - ))); - } - /** * Given explicit runtime construction, config snapshots reject values above the delay upper bound. */ @@ -121,7 +101,7 @@ public void testPartialDynamicUpdate() { assertFalse(config.enabled()); assertEquals(1000L, config.maxDelayMs()); - config.update(Map.of(AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG, "250")); + config.update(Map.of(AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG, 250L)); assertFalse(config.enabled()); assertEquals(250L, config.maxDelayMs()); } diff --git a/core/src/test/java/kafka/automq/retrystorm/RetryStormBackoffManagerTest.java b/core/src/test/java/kafka/automq/retrystorm/RetryStormBackoffManagerTest.java index 7da60ccead..abf608f820 100644 --- a/core/src/test/java/kafka/automq/retrystorm/RetryStormBackoffManagerTest.java +++ b/core/src/test/java/kafka/automq/retrystorm/RetryStormBackoffManagerTest.java @@ -22,7 +22,6 @@ import kafka.automq.AutoMQConfig; import kafka.server.retrystorm.RetryStormBackoffLogger; -import org.apache.kafka.common.config.ConfigException; import org.apache.kafka.common.protocol.ApiKeys; import org.junit.jupiter.api.Tag; @@ -34,7 +33,6 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertFalse; -import static org.junit.jupiter.api.Assertions.assertThrows; /** * Covers manager integration with Kafka dynamic broker config callbacks. @@ -59,18 +57,6 @@ public void testReconfigureUpdatesRuntimeConfig() { assertEquals(250L, config.maxDelayMs()); } - /** - * Given an invalid dynamic update, manager validation rejects it before runtime state changes. - */ - @Test - public void testValidateRejectsInvalidReconfiguration() { - RetryStormBackoffManager manager = new RetryStormBackoffManager(new RetryStormBackoffConfig(true, 1000L)); - - assertThrows(ConfigException.class, () -> manager.validateReconfiguration(Map.of( - AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG, -1L - ))); - } - /** * Given many old delaying states, one periodic logging scan emits no more than ten snapshots. */ diff --git a/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala index 128b03bcd5..0e7b1039f4 100755 --- a/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala +++ b/core/src/test/scala/unit/kafka/server/DynamicBrokerConfigTest.scala @@ -23,6 +23,7 @@ import java.util.concurrent.CompletionStage import java.util.concurrent.atomic.AtomicReference import kafka.controller.KafkaController import kafka.automq.AutoMQConfig +import kafka.automq.metadata.KafkaGoMetadataResponseRewriter import kafka.log.LogManager import kafka.log.remote.RemoteLogManager import kafka.network.{DataPlaneAcceptor, SocketServer} @@ -66,6 +67,15 @@ class DynamicBrokerConfigTest { assertTrue(DynamicBrokerConfig.AllDynamicConfigs.contains(AutoMQConfig.RETRY_STORM_BACKOFF_MAX_DELAY_MS_CONFIG)) } + /** Given kafka-go metadata compatibility is configurable, the registry exposes its key for AdminClient updates. */ + @Test + @Tag("S3Unit") + def testKafkaGoMetadataCompatibilityConfigIsDynamic(): Unit = { + assertTrue(KafkaGoMetadataResponseRewriter.RECONFIGURABLE_CONFIGS.asScala + .forall(DynamicBrokerConfig.AllDynamicConfigs.contains)) + assertTrue(KafkaConfig(TestUtils.createBrokerConfig(0, null)).kafkaGoMetadataCompatibilityEnabled) + } + @Test def testConfigUpdate(): Unit = { val props = TestUtils.createBrokerConfig(0, TestUtils.MockZkConnect, port = 8181)