void addInputChannel(RemoteInputChannel listener) {
    checkState(!channelError.get(), "There has been an error in the channel.");

    if (!inputChannels.containsKey(listener.getInputChannelId())) {
      inputChannels.put(listener.getInputChannelId(), listener);
    }
  }
  private boolean decodeMsg(Object msg) throws Throwable {
    final Class<?> msgClazz = msg.getClass();

    // ---- Buffer --------------------------------------------------------
    if (msgClazz == NettyMessage.BufferResponse.class) {
      NettyMessage.BufferResponse bufferOrEvent = (NettyMessage.BufferResponse) msg;

      RemoteInputChannel inputChannel = inputChannels.get(bufferOrEvent.receiverId);
      if (inputChannel == null) {
        bufferOrEvent.releaseBuffer();

        cancelRequestFor(bufferOrEvent.receiverId);

        return true;
      }

      return decodeBufferOrEvent(inputChannel, bufferOrEvent);
    }
    // ---- Error ---------------------------------------------------------
    else if (msgClazz == NettyMessage.ErrorResponse.class) {
      NettyMessage.ErrorResponse error = (NettyMessage.ErrorResponse) msg;

      SocketAddress remoteAddr = ctx.channel().remoteAddress();

      if (error.isFatalError()) {
        notifyAllChannelsOfErrorAndClose(
            new RemoteTransportException(
                "Fatal error at remote task manager '" + remoteAddr + "'.",
                remoteAddr,
                error.cause));
      } else {
        RemoteInputChannel inputChannel = inputChannels.get(error.receiverId);

        if (inputChannel != null) {
          if (error.cause.getClass() == PartitionNotFoundException.class) {
            inputChannel.onFailedPartitionRequest();
          } else {
            inputChannel.onError(
                new RemoteTransportException(
                    "Error at remote task manager '" + remoteAddr + "'.", remoteAddr, error.cause));
          }
        }
      }
    } else {
      throw new IllegalStateException("Received unknown message from producer: " + msg.getClass());
    }

    return true;
  }
    /**
     * Continues the decoding of a staged buffer after a buffer has become available again.
     *
     * <p>This task is executed by the network I/O thread.
     */
    @Override
    public void run() {
      boolean success = false;

      Buffer buffer = null;

      try {
        if ((buffer = availableBuffer.getAndSet(null)) == null) {
          throw new IllegalStateException("Running buffer availability task w/o a buffer.");
        }

        buffer.setSize(stagedBufferResponse.getSize());

        stagedBufferResponse.getNettyBuffer().readBytes(buffer.getNioBuffer());
        stagedBufferResponse.releaseBuffer();

        RemoteInputChannel inputChannel = inputChannels.get(stagedBufferResponse.receiverId);

        if (inputChannel != null) {
          inputChannel.onBuffer(buffer, stagedBufferResponse.sequenceNumber);

          success = true;
        } else {
          cancelRequestFor(stagedBufferResponse.receiverId);
        }

        stagedBufferResponse = null;

        if (stagedMessages.isEmpty()) {
          ctx.channel().config().setAutoRead(true);
          ctx.channel().read();
        } else {
          ctx.channel().eventLoop().execute(stagedMessagesHandler);
        }
      } catch (Throwable t) {
        notifyAllChannelsOfErrorAndClose(t);
      } finally {
        if (!success) {
          if (buffer != null) {
            buffer.recycle();
          }
        }
      }
    }
  private void notifyAllChannelsOfErrorAndClose(Throwable cause) {
    if (channelError.compareAndSet(false, true)) {
      try {
        for (RemoteInputChannel inputChannel : inputChannels.values()) {
          inputChannel.onError(cause);
        }
      } catch (Throwable t) {
        // We can only swallow the Exception at this point. :(
        LOG.warn(
            "An Exception was thrown during error notification of a " + "remote input channel.", t);
      } finally {
        inputChannels.clear();

        if (ctx != null) {
          ctx.close();
        }
      }
    }
  }
 void removeInputChannel(RemoteInputChannel listener) {
   inputChannels.remove(listener.getInputChannelId());
 }
  private boolean decodeBufferOrEvent(
      RemoteInputChannel inputChannel, NettyMessage.BufferResponse bufferOrEvent) throws Throwable {
    boolean releaseNettyBuffer = true;

    try {
      if (bufferOrEvent.isBuffer()) {
        // ---- Buffer ------------------------------------------------

        // Early return for empty buffers. Otherwise Netty's readBytes() throws an
        // IndexOutOfBoundsException.
        if (bufferOrEvent.getSize() == 0) {
          inputChannel.onEmptyBuffer(bufferOrEvent.sequenceNumber);
          return true;
        }

        BufferProvider bufferProvider = inputChannel.getBufferProvider();

        if (bufferProvider == null) {

          cancelRequestFor(bufferOrEvent.receiverId);

          return false; // receiver has been cancelled/failed
        }

        while (true) {
          Buffer buffer = bufferProvider.requestBuffer();

          if (buffer != null) {
            buffer.setSize(bufferOrEvent.getSize());
            bufferOrEvent.getNettyBuffer().readBytes(buffer.getNioBuffer());

            inputChannel.onBuffer(buffer, bufferOrEvent.sequenceNumber);

            return true;
          } else if (bufferListener.waitForBuffer(bufferProvider, bufferOrEvent)) {
            releaseNettyBuffer = false;

            return false;
          } else if (bufferProvider.isDestroyed()) {
            return false;
          }
        }
      } else {
        // ---- Event -------------------------------------------------
        // TODO We can just keep the serialized data in the Netty buffer and release it later at the
        // reader
        byte[] byteArray = new byte[bufferOrEvent.getSize()];
        bufferOrEvent.getNettyBuffer().readBytes(byteArray);

        Buffer buffer =
            new Buffer(new MemorySegment(byteArray), FreeingBufferRecycler.INSTANCE, false);

        inputChannel.onBuffer(buffer, bufferOrEvent.sequenceNumber);

        return true;
      }
    } finally {
      if (releaseNettyBuffer) {
        bufferOrEvent.releaseBuffer();
      }
    }
  }