io.streamnative.pulsar.handlers.kop.proxy.ConnectionFactory 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 io.netty.bootstrap.Bootstrap;
import io.netty.channel.Channel;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelInitializer;
import io.netty.channel.ChannelOption;
import io.netty.channel.EventLoopGroup;
import io.netty.channel.socket.SocketChannel;
import io.netty.handler.codec.LengthFieldBasedFrameDecoder;
import io.netty.handler.codec.LengthFieldPrepender;
import io.netty.resolver.AddressResolver;
import io.netty.resolver.dns.DnsAddressResolverGroup;
import io.netty.util.concurrent.Future;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.util.List;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.atomic.AtomicInteger;
import lombok.extern.slf4j.Slf4j;
import org.apache.pulsar.common.allocator.PulsarByteBufAllocator;
import org.apache.pulsar.common.util.netty.EventLoopUtil;
@Slf4j
class ConnectionFactory {
private final AtomicInteger id = new AtomicInteger(0);
private final List addresses;
private final int connectTimeoutMs;
private final EventLoopGroup eventLoopGroup;
private final AddressResolver addressResolver;
ConnectionFactory(final KafkaProxyConfiguration config, final EventLoopGroup eventLoopGroup,
final DnsAddressResolverGroup dnsAddressResolverGroup) {
this.addresses = config.getKafkaBootstrapServers();
this.connectTimeoutMs = config.getBrokerProxyConnectTimeoutMs();
this.eventLoopGroup = eventLoopGroup;
this.addressResolver = dnsAddressResolverGroup.getResolver(eventLoopGroup.next());
}
ConnectionToBroker getAnyConnection() throws IOException {
for (int i = 0; i < addresses.size(); i++) {
final var address = addresses.get(id.getAndIncrement() % addresses.size());
try {
return getConnection(address);
} catch (IOException e) {
if (i == addresses.size() - 1) {
throw e;
}
}
}
throw new IOException("Unknown error");
}
ConnectionToBroker getConnection(final InetSocketAddress unresolvedAddress) throws IOException {
final var addresses = waitFuture(addressResolver.resolveAll(unresolvedAddress),
"Resolve " + unresolvedAddress);
final var bootstrap = new Bootstrap();
bootstrap.group(eventLoopGroup);
bootstrap.channel(EventLoopUtil.getClientSocketChannelClass(eventLoopGroup));
bootstrap.option(ChannelOption.CONNECT_TIMEOUT_MILLIS, connectTimeoutMs);
bootstrap.option(ChannelOption.TCP_NODELAY, true);
bootstrap.option(ChannelOption.ALLOCATOR, PulsarByteBufAllocator.DEFAULT);
bootstrap.handler(new ChannelInitializer() {
@Override
protected void initChannel(SocketChannel ch) throws Exception {
ch.pipeline().addLast(new LengthFieldPrepender(4));
ch.pipeline().addLast("frameDecoder",
new LengthFieldBasedFrameDecoder(100 * 1024 * 1024/* 100 MB */, 0, 4, 0, 4));
ch.pipeline().addLast("handler", new ConnectionToBroker(unresolvedAddress));
}
});
final var registeredChannel = waitChannelFuture(bootstrap.register(), "Register channel");
for (int i = 0; i < addresses.size(); i++) {
final var address = addresses.get(i);
try {
final var channel = waitChannelFuture(registeredChannel.connect(address), "Connect " + address);
final var cnx = (ConnectionToBroker) channel.pipeline().get("handler");
if (cnx == null) {
throw new IOException("null handler in the pipeline of " + channel);
}
return cnx;
} catch (IOException e) {
if (i == addresses.size() - 1) {
throw e;
}
}
}
throw new IOException("Unknown error"); // it should never reach here
}
private static T waitFuture(final Future future, final String msg) throws IOException {
try {
if (log.isDebugEnabled()) {
future.addListener(innerFuture -> {
if (innerFuture.isSuccess()) {
log.debug("{} succeeded", msg);
} else {
log.debug("{} failed: {}", msg, innerFuture.cause());
}
});
}
return future.get();
} catch (ExecutionException e) {
throw new IOException(e.getCause());
} catch (InterruptedException e) {
throw new IOException(Thread.currentThread().getName() + " is interrupted");
}
}
private static Channel waitChannelFuture(final ChannelFuture channelFuture, final String msg) throws IOException {
waitFuture(channelFuture, msg);
return channelFuture.channel();
}
}