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

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

Go to download

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 static org.apache.kafka.common.protocol.ApiKeys.API_VERSIONS;

import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.util.ReferenceCounted;
import java.io.IOException;
import java.net.SocketAddress;
import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.function.Consumer;
import java.util.function.Function;
import lombok.Getter;
import lombok.Setter;
import org.apache.commons.lang3.tuple.Pair;
import org.apache.kafka.common.protocol.ApiKeys;
import org.apache.kafka.common.requests.AbstractRequest;
import org.apache.kafka.common.requests.AbstractResponse;
import org.apache.kafka.common.requests.AddOffsetsToTxnRequest;
import org.apache.kafka.common.requests.AddPartitionsToTxnRequest;
import org.apache.kafka.common.requests.ApiVersionsRequest;
import org.apache.kafka.common.requests.EndTxnRequest;
import org.apache.kafka.common.requests.HeartbeatRequest;
import org.apache.kafka.common.requests.InitProducerIdRequest;
import org.apache.kafka.common.requests.JoinGroupRequest;
import org.apache.kafka.common.requests.KopResponseUtils;
import org.apache.kafka.common.requests.LeaveGroupRequest;
import org.apache.kafka.common.requests.OffsetCommitRequest;
import org.apache.kafka.common.requests.OffsetDeleteRequest;
import org.apache.kafka.common.requests.OffsetFetchRequest;
import org.apache.kafka.common.requests.RequestHeader;
import org.apache.kafka.common.requests.SyncGroupRequest;
import org.apache.kafka.common.requests.TxnOffsetCommitRequest;

public class InflightRequest {

    private final CompletableFuture responseFuture = new CompletableFuture<>();
    @Getter
    private final RequestHeader header;
    // The response can be a ByteBuf or an AbstractResponse implementation
    @Getter
    private final AbstractRequest request;
    @Getter
    private final SocketAddress remoteAddress;
    private final ByteBuf buf;
    @Setter
    private Function responseMapper = __ -> __;
    @Getter
    @Setter
    private boolean skipParsingResponse = false;

    public InflightRequest(final ByteBuf buf, final SocketAddress remoteAddress) {
        this(buf, remoteAddress, true);
    }

    public InflightRequest(final ByteBuf buf, final SocketAddress remoteAddress, final boolean parseRequest) {
        this.buf = buf.retain();
        this.remoteAddress = remoteAddress;
        final var nio = buf.nioBuffer();
        this.header = RequestHeader.parse(nio);
        switch (header.apiKey()) {
            case API_VERSIONS -> {
                if (API_VERSIONS.isVersionSupported(header.apiVersion())) {
                    this.request = AbstractRequest.parseRequest(header.apiKey(), header.apiVersion(), nio).request;
                } else {
                    this.request = new ApiVersionsRequest.Builder(header.apiVersion()).build();
                }
            }
            case FIND_COORDINATOR, PRODUCE, LIST_OFFSETS, FETCH -> {
                if (parseRequest) {
                    this.request = AbstractRequest.parseRequest(header.apiKey(), header.apiVersion(), nio).request;
                } else {
                    this.request = null;
                }
            }
            case METADATA -> this.request = null;
            default -> this.request = AbstractRequest.parseRequest(header.apiKey(), header.apiVersion(), nio).request;
        }
    }

    @SuppressWarnings("unchecked")
    public ByteBuf toResponseBuf() {
        final var response = responseFuture.join();
        if (response instanceof ByteBuf) {
            return (ByteBuf) response;
        } else if (response instanceof AbstractResponse) {
            // Lowering Client API_VERSION request to the oldest API_VERSION KoP supports, this is to make \
            // Kafka-Clients 2.4.x and above compatible and prevent KoP from panicking \
            // when it comes across a higher API_VERSION.
            short apiVersion = header.apiVersion();
            if (header.apiKey() == API_VERSIONS){
                if (!ApiKeys.API_VERSIONS.isVersionSupported(apiVersion)) {
                    apiVersion = ApiKeys.API_VERSIONS.oldestVersion();
                }
            }
            return KopResponseUtils.serializeResponse(apiVersion, header.toResponseHeader(),
                    (AbstractResponse) response);
        } else {
            // Fetch request that was split into multiple requests to different brokers
            final var pair = (Pair>) response;
            try {
                return KopResponseUtils.serializeResponse(header.apiVersion(), header.toResponseHeader(),
                        pair.getLeft());
            } finally {
                pair.getValue().forEach(ReferenceCounted::release);
            }
        }
    }

    public void sendToChannel(final Channel channel) {
        channel.writeAndFlush(buf);
    }

    public void registerCallback(final Runnable callback, final Executor executor) {
        responseFuture.whenCompleteAsync((__, ___) -> callback.run(), executor);
    }

    public void complete(final Object response) {
        responseFuture.complete(responseMapper.apply(response));
    }

    public void fail(final Throwable e) {
        responseFuture.completeExceptionally(e);
    }

    public boolean hasReceivedResponse() {
        return responseFuture.isDone();
    }

    public boolean hasFailed(final Consumer throwableConsumer) {
        if (!responseFuture.isCompletedExceptionally()) {
            return false;
        }
        responseFuture.exceptionally(e -> {
            throwableConsumer.accept(e);
            return null;
        });
        return true;
    }

    public ByteBuf getRetainedBuffer() {
        return buf.retain();
    }

    public Object waitForResponse() throws IOException {
        try {
            return responseFuture.get();
        } catch (ExecutionException e) {
            throw new IOException(e.getCause());
        } catch (InterruptedException e) {
            throw new IOException(this + " is interrupted");
        }
    }

    @SuppressWarnings("unchecked")
    public  CompletableFuture getResponseFuture() {
        return (CompletableFuture) responseFuture;
    }

    public String groupId() {
        return switch (header.apiKey()) {
            case JOIN_GROUP -> ((JoinGroupRequest) request).data().groupId();
            case SYNC_GROUP -> ((SyncGroupRequest) request).data().groupId();
            case LEAVE_GROUP -> ((LeaveGroupRequest) request).data().groupId();
            case OFFSET_FETCH -> getGroupIdForOffsetFetch((OffsetFetchRequest) request);
            case OFFSET_COMMIT -> ((OffsetCommitRequest) request).data().groupId();
            case HEARTBEAT -> ((HeartbeatRequest) request).data().groupId();
            case TXN_OFFSET_COMMIT -> ((TxnOffsetCommitRequest) request).data().groupId();
            case OFFSET_DELETE -> ((OffsetDeleteRequest) request).data().groupId();
            default -> throw new IllegalStateException("Cannot call groupId() for apiKey=" + header.apiKey());
        };
    }

    public String txnId() {
        return switch (header.apiKey()) {
            case INIT_PRODUCER_ID -> ((InitProducerIdRequest) request).data().transactionalId();
            case ADD_PARTITIONS_TO_TXN -> ((AddPartitionsToTxnRequest) request).data().transactionalId();
            case ADD_OFFSETS_TO_TXN -> ((AddOffsetsToTxnRequest) request).data().transactionalId();
            case END_TXN -> ((EndTxnRequest) request).data().transactionalId();
            default -> throw new IllegalStateException("Cannot call txnId() for apiKey=" + header.apiKey());
        };
    }

    private String getGroupIdForOffsetFetch(final OffsetFetchRequest request) {
        if (request.version() >= 8) {
            // TODO: support multiple groups for OffsetFetch request v8 or later
            if (request.groupIds().size() > 1) {
                throw new IllegalStateException("KoP proxy does not support v8 OffsetFetch with multiple groups");
            }
            return request.groupIds().get(0);
        } else {
            return request.groupId();
        }
    }

    @Override
    public String toString() {
        if (request != null) {
            return String.format("InflightRequest(header=%s, request=%s, remoteAddress=%s)",
                    header, request, remoteAddress);
        } else {
            return String.format("InflightRequest(header=%s, remoteAddress=%s)", header, remoteAddress);
        }
    }
}
    

    

    

            

    

            



    
    






    
© 2015 - 2024 Weber Informatics LLC | Privacy Policy