io.streamnative.pulsar.handlers.kop.proxy.BrokerConnectionGroup Maven / Gradle / Ivy
/**
* Copyright (c) 2019 - 2024 StreamNative, Inc.. All Rights Reserved.
*/
/**
* Licensed 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 io.streamnative.pulsar.handlers.kop.proxy;
import static org.apache.kafka.common.requests.FindCoordinatorRequest.CoordinatorType;
import com.google.common.annotations.VisibleForTesting;
import io.netty.buffer.ByteBuf;
import io.netty.channel.ChannelHandlerContext;
import io.netty.util.ReferenceCounted;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CopyOnWriteArrayList;
import lombok.RequiredArgsConstructor;
import lombok.Setter;
import org.apache.kafka.common.message.FindCoordinatorRequestData;
import org.apache.kafka.common.protocol.ApiKeys;
import org.apache.kafka.common.requests.FindCoordinatorRequest;
import org.apache.kafka.common.requests.FindCoordinatorResponse;
import org.apache.kafka.common.requests.KopResponseUtils;
import org.apache.kafka.common.requests.RequestHeader;
/**
* This class maintains connections of a Kafka client that could establish:
* 1. The connection for METADATA requests
* 2. The connection for group coordinator related requests.
* 3. The connection to leader brokers for PRODUCE and FETCH requests.
*/
@RequiredArgsConstructor
public class BrokerConnectionGroup {
private static final InetSocketAddress DUMMY_ADDRESS = InetSocketAddress.createUnresolved("localhost", 19092);
private final ConnectionFactory connectionFactory;
@VisibleForTesting
final Map leaderBrokers = new ConcurrentHashMap<>();
private final List saslRequestBuffers = new CopyOnWriteArrayList<>();
@VisibleForTesting
volatile ConnectionToBroker metadataBroker;
volatile ConnectionToBroker groupCoordinator;
volatile ConnectionToBroker transactionCoordinator;
// This connection is only used to verify if the SASL_AUTHENTICATE requests succeed
@VisibleForTesting
volatile ConnectionToBroker authenticateBroker;
@Setter
private ChannelHandlerContext clientChannel = null;
public ConnectionToBroker getLeader(final InetSocketAddress leader) throws IOException {
var connectionToBroker = leaderBrokers.get(leader);
if (connectionToBroker == null) {
connectionToBroker = connectionFactory.getConnection(leader).withClientChannel(clientChannel)
.withDisconnectCallback(() -> leaderBrokers.remove(leader))
.forwardRequestsAndWait(createSaslRequests());
leaderBrokers.put(leader, connectionToBroker);
}
return connectionToBroker;
}
public ConnectionToBroker getMetadataBroker() throws IOException {
if (metadataBroker != null) {
return metadataBroker;
}
synchronized (this) {
if (metadataBroker == null) {
metadataBroker = connectionFactory.getAnyConnection().withClientChannel(clientChannel)
.withDisconnectCallback(() -> metadataBroker = null)
.forwardRequestsAndWait(createSaslRequests());
}
}
return metadataBroker;
}
public ConnectionToBroker getGroupCoordinator(final String groupId) throws IOException {
if (groupCoordinator != null) {
return groupCoordinator;
}
synchronized (this) {
if (groupCoordinator == null) {
final var address = findCoordinator(groupId, CoordinatorType.GROUP);
groupCoordinator = connectionFactory.getConnection(address).withClientChannel(clientChannel)
.withDisconnectCallback(() -> groupCoordinator = null)
.forwardRequestsAndWait(createSaslRequests());
}
}
return groupCoordinator;
}
public ConnectionToBroker getTransactionCoordinator(final String txnId) throws IOException {
if (txnId == null) {
// It's an INIT_PRODUCER_ID request from the idempotent producer
return getMetadataBroker();
}
if (transactionCoordinator != null) {
return transactionCoordinator;
}
synchronized (this) {
if (transactionCoordinator == null) {
final var address = findCoordinator(txnId, CoordinatorType.TRANSACTION);
transactionCoordinator = connectionFactory.getConnection(address).withClientChannel(clientChannel)
.withDisconnectCallback(() -> transactionCoordinator = null)
.forwardRequestsAndWait(createSaslRequests());
}
}
return transactionCoordinator;
}
private InetSocketAddress findCoordinator(final String key, final CoordinatorType keyType) throws IOException {
final var request = new FindCoordinatorRequest.Builder(new FindCoordinatorRequestData()
.setKey(key).setKeyType(keyType.id())).build((short) 3);
final var buf = KopResponseUtils.serializeRequestToPooledBuffer(new RequestHeader(ApiKeys.FIND_COORDINATOR,
(short) 3, "kop-proxy", 0), request);
final var inflightRequest = new InflightRequest(buf, DUMMY_ADDRESS);
getMetadataBroker().forwardRequest(inflightRequest);
final var node = ((FindCoordinatorResponse) inflightRequest.waitForResponse()).node();
return InetSocketAddress.createUnresolved(node.host(), node.port());
}
@VisibleForTesting
ConnectionToBroker getAuthenticateBroker() throws IOException {
if (authenticateBroker != null) {
return authenticateBroker;
}
synchronized (this) {
if (authenticateBroker == null) {
authenticateBroker = connectionFactory.getAnyConnection()
.withDisconnectCallback(() -> authenticateBroker = null);
}
}
return authenticateBroker;
}
public void authenticate(final InflightRequest inflightRequest) throws IOException {
getAuthenticateBroker().forwardRequest(inflightRequest);
saslRequestBuffers.add(inflightRequest.getRetainedBuffer());
}
public void addSaslRequestBuffer(final InflightRequest inflightRequest) {
saslRequestBuffers.add(inflightRequest.getRetainedBuffer());
}
public void close() {
saslRequestBuffers.forEach(ReferenceCounted::release);
saslRequestBuffers.clear();
clientChannel = null;
if (metadataBroker != null) {
metadataBroker.disconnectBroker();
metadataBroker = null;
}
if (groupCoordinator != null) {
groupCoordinator.disconnectBroker();
groupCoordinator = null;
}
if (authenticateBroker != null) {
authenticateBroker.disconnectBroker();
authenticateBroker = null;
}
leaderBrokers.values().forEach(ConnectionToBroker::disconnectBroker);
leaderBrokers.clear();
}
private List createSaslRequests() {
final var address = (clientChannel != null) ? clientChannel.channel().remoteAddress()
: InetSocketAddress.createUnresolved("localhost", 65535); // fake address
return saslRequestBuffers.stream().map(buffer -> new InflightRequest(buffer, address)).toList();
}
}