bbejeck commented on code in PR #23438: URL: https://github.com/apache/kafka/pull/23438#discussion_r4042590340
########## clients/src/testFixtures/java/org/apache/kafka/test/faultproxy/KafkaProtocolFaultProxy.java: ########## @@ -0,0 +1,450 @@ +/* + * 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.kafka.test.faultproxy; + +import org.apache.kafka.common.message.FetchResponseData; +import org.apache.kafka.common.message.ProduceResponseData; +import org.apache.kafka.common.message.ResponseHeaderData; +import org.apache.kafka.common.message.TxnOffsetCommitResponseData; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.AbstractResponse; +import org.apache.kafka.common.requests.AddOffsetsToTxnResponse; +import org.apache.kafka.common.requests.EndTxnResponse; +import org.apache.kafka.common.requests.FetchResponse; +import org.apache.kafka.common.requests.FindCoordinatorResponse; +import org.apache.kafka.common.requests.InitProducerIdResponse; +import org.apache.kafka.common.requests.MetadataResponse; +import org.apache.kafka.common.requests.RequestHeader; +import org.apache.kafka.common.requests.RequestUtils; +import org.apache.kafka.common.requests.TxnOffsetCommitResponse; +import org.apache.kafka.common.utils.Utils; + +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; + +import java.io.DataInputStream; +import java.io.DataOutputStream; +import java.net.ServerSocket; +import java.net.Socket; +import java.nio.ByteBuffer; +import java.time.Duration; +import java.util.EnumMap; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.BiConsumer; + +/** + * A lightweight, client-agnostic Kafka wire-protocol fault-injection proxy for fast integration tests. + * + * <p>Sit it in front of a real (embedded) broker and point any client's {@code bootstrap.servers} at it; + * it decodes requests/responses with Kafka's own protocol classes ({@link RequestHeader}, + * {@link AbstractResponse}, {@link RequestUtils#serialize}) — so it is correct across every wire version, + * including flexible/tagged-field ones, with no hand-rolled byte offsets. + * + * <h2>Usage</h2> + * <pre>{@code + * try (var broker = new EmbeddedKafkaCluster(1)) { + * broker.start(); + * try (var proxy = KafkaProtocolFaultProxy.inFrontOf(broker.bootstrapServers())) { + * // point clients here: + * props.put(BOOTSTRAP_SERVERS_CONFIG, proxy.bootstrapServers()); + * + * proxy.injectError(ApiKeys.END_TXN, Errors.CONCURRENT_TRANSACTIONS).once(); + * proxy.injectError(ApiKeys.PRODUCE, Errors.NOT_ENOUGH_REPLICAS).onCall(2); + * proxy.disconnectOn(ApiKeys.END_TXN).once(); // the EOS "commit gap" + * proxy.delayOn(ApiKeys.FETCH, Duration.ofSeconds(2)).everyTime(); // slow broker + * } + * } + * }</pre> + * + * <p>Routing is transparent: the proxy rewrites {@code Metadata}/{@code FindCoordinator} responses so every + * advertised address points back at itself, so a single-broker embedded cluster needs no special config + * (its own ephemeral port is discovered from {@code bootstrapServers()}). + * + * <p>Determinism: {@code once()}/{@code onCall(n)}/{@code times(n)} are deterministic and safe for + * assertions; {@code withProbability(p)} is chaos-mode only. The proxy never closes sockets unless a + * {@code disconnectOn(...)} rule fires, so it is not itself a source of flakiness. + */ +public final class KafkaProtocolFaultProxy implements AutoCloseable { + + private static final Logger LOG = LoggerFactory.getLogger(KafkaProtocolFaultProxy.class); + + /** + * Per-API setter that stamps an injected {@link Errors} onto a decoded response. Only APIs registered + * here support {@code injectError(...)}; anything else fails fast at rule registration. + */ + private static final Map<ApiKeys, BiConsumer<AbstractResponse, Errors>> ERROR_SETTERS = new EnumMap<>(ApiKeys.class); + static { + ERROR_SETTERS.put(ApiKeys.END_TXN, (r, e) -> ((EndTxnResponse) r).data().setErrorCode(e.code())); + ERROR_SETTERS.put(ApiKeys.INIT_PRODUCER_ID, (r, e) -> ((InitProducerIdResponse) r).data().setErrorCode(e.code())); + ERROR_SETTERS.put(ApiKeys.ADD_OFFSETS_TO_TXN, (r, e) -> ((AddOffsetsToTxnResponse) r).data().setErrorCode(e.code())); + // TxnOffsetCommit carries the consumed offsets into the transaction. Under EOS-v2 / transactions V2 + // (KIP-890) the client sends this directly (AddOffsetsToTxn is skipped -- see TransactionManager + // #sendOffsetsToTransaction), so this is THE offset-commit-into-txn RPC to fault for KIP-1035. The + // response carries per-partition error codes, so stamp every partition of every topic. + ERROR_SETTERS.put(ApiKeys.TXN_OFFSET_COMMIT, (r, e) -> { + final TxnOffsetCommitResponseData data = ((TxnOffsetCommitResponse) r).data(); + data.topics().forEach(topic -> + topic.partitions().forEach(p -> p.setErrorCode(e.code()))); + }); + ERROR_SETTERS.put(ApiKeys.PRODUCE, (r, e) -> { + final ProduceResponseData data = ((org.apache.kafka.common.requests.ProduceResponse) r).data(); + data.responses().forEach(topic -> + topic.partitionResponses().forEach(p -> p.setErrorCode(e.code()))); + }); + // FETCH stamps the error on every partition of the response. Because a fetch fault is almost always + // scoped with forClient("restore") (or another clientId), this hits only the targeted consumer's + // fetches — e.g. inject OFFSET_OUT_OF_RANGE on the restore consumer to exercise the restore path. + ERROR_SETTERS.put(ApiKeys.FETCH, (r, e) -> { + final FetchResponseData data = ((FetchResponse) r).data(); + data.responses().forEach(topic -> + topic.partitions().forEach(p -> p.setErrorCode(e.code()))); + }); + } + + private final String targetHost; + private final int targetPort; + private final ExecutorService threadPool = Executors.newCachedThreadPool(r -> { + final Thread t = new Thread(r, "kafka-fault-proxy"); + t.setDaemon(true); + return t; + }); + private final AtomicBoolean running = new AtomicBoolean(false); + private final CopyOnWriteArrayList<FaultRule> rules = new CopyOnWriteArrayList<>(); + private final Set<String> blackholedClients = ConcurrentHashMap.newKeySet(); + // Live connections, so close() can force sockets shut and unblock the pump reads that own them. + private final Set<Connection> connections = ConcurrentHashMap.newKeySet(); + private ServerSocket serverSocket; + private volatile String proxyHost; + private volatile int proxyPort; + + private KafkaProtocolFaultProxy(final String targetBootstrap) { + final String hostPort = targetBootstrap.split(",")[0].trim(); Review Comment: No I don't think so becuase in integration tests we specify any number of brokers (between 1 and 3 from my experience) I'll drill into this a little more and file a follow-up PR if need be. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
