All Downloads are FREE. Search and download functionalities are using the official Maven repository.

io.streamnative.pulsar.handlers.kop.proxy.ConnectionFactory Maven / Gradle / Ivy

The newest version!
/**
 * 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();
    }

}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy