io.streamnative.pulsar.handlers.kop.proxy.ConnectionToBroker 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.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);
}
}
}