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

Go to download
Show more of this group Show more artifacts with this name
Show all versions of pulsar-kafka-proxy Show documentation
Kafka on Pulsar Proxy
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.buffer.ByteBuf;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.util.ReferenceCountUtil;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicBoolean;
import lombok.Getter;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.common.requests.AbstractResponse;

@RequiredArgsConstructor
@Slf4j
public class ConnectionToBroker extends ChannelInboundHandlerAdapter {

    private final Map pendingRequests = new ConcurrentHashMap<>();

    // forwardRequest might be called in another thread, so here use volatile to ensure the safe access to ctx
    private final AtomicBoolean closed = new AtomicBoolean(false);
    @Getter
    private final InetSocketAddress address;
    private volatile ChannelHandlerContext ctx = null;
    private volatile ChannelHandlerContext clientChannel = null;
    private volatile Runnable disconnectCallback = null;

    @Override
    public void channelRegistered(ChannelHandlerContext ctx) throws Exception {
        super.channelRegistered(ctx);
        this.ctx = ctx;
        log.info("[{}] Connection to broker is registered", ctx);
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) throws Exception {
        log.error("[{}] Unexpected error in ConnectionToBroker", ctx, cause);
        close();
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg) throws Exception {
        final var buf = (ByteBuf) msg;
        try {
            final var readerIndex = buf.readerIndex();
            final var correlationId = buf.getInt(0);
            final var inflightRequest = pendingRequests.remove(correlationId);
            if (inflightRequest == null) {
                log.warn("[{}] Correlation id {} is not pending", ctx, correlationId);
                close();
                return;
            }
            buf.readerIndex(readerIndex);

            if (inflightRequest.isSkipParsingResponse()) {
                if (log.isDebugEnabled()) {
                    log.debug("[{}] Received response of {} from broker", ctx, inflightRequest.getHeader());
                }
                inflightRequest.complete(buf.retain());
                return;
            }

            final ByteBuffer buffer = buf.nioBuffer();
            final var response = AbstractResponse.parseResponse(buffer, inflightRequest.getHeader());
            if (log.isDebugEnabled()) {
                log.debug("[{}] Received response {} from broker", ctx, response);
            }
            inflightRequest.complete(response);
        } catch (Throwable throwable) {
            log.error("[{}] Unexpected error when handling responses from broker", ctx, throwable);
            close();
        } finally {
            ReferenceCountUtil.safeRelease(buf);
        }
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        log.info("[{}] Connection to broker is inactive", ctx);
        close();
        this.ctx = null;
    }

    void disconnectBroker() {
        if (ctx != null) {
            ctx.close();
        }
    }

    private void close() {
        if (closed.compareAndSet(false, true)) {
            if (clientChannel != null) {
                log.info("[{}] Close connection to client {}", ctx, clientChannel);
                clientChannel.close();
            }
            disconnectBroker();
            pendingRequests.values().forEach(inflightRequest -> inflightRequest.fail(
                    new ConnectError("Connection is closed")));
            pendingRequests.clear();
            if (disconnectCallback != null) {
                disconnectCallback.run();
            }
        }
    }

    void forwardRequest(final InflightRequest inflightRequest) {
        forwardRequest(inflightRequest, true);
    }

    void forwardRequest(final InflightRequest inflightRequest, final boolean cache) {
        final var ctx = this.ctx;
        if (ctx == null) {
            log.error("Channel is inactive when forwarding request {}", inflightRequest.getHeader());
            inflightRequest.fail(new ConnectError("Channel is not registered"));
            return;
        }
        if (closed.get()) {
            inflightRequest.fail(new ConnectError("Connection is closed"));
            return;
        }
        final var correlationId = inflightRequest.getHeader().correlationId();
        if (cache && pendingRequests.putIfAbsent(correlationId, inflightRequest) != null) {
            log.error("[{}] Received request with the same correlation id {}", ctx, correlationId);
            inflightRequest.fail(new ConnectError("Duplicated correlation id " + correlationId));
            return;
        }
        inflightRequest.sendToChannel(ctx.channel());
    }

    ConnectionToBroker withDisconnectCallback(final Runnable callback) {
        this.disconnectCallback = callback;
        return this;
    }

    ConnectionToBroker withClientChannel(final ChannelHandlerContext clientChannel) {
        this.clientChannel = clientChannel;
        return this;
    }

    ConnectionToBroker forwardRequestsAndWait(final List requests) throws IOException {
        requests.forEach(this::forwardRequest);
        for (var request : requests) {
            request.waitForResponse();
        }
        return this;
    }

    static class ConnectError extends RuntimeException {

        public ConnectError(final String msg) {
            super(msg);
        }
    }
}