content: /*
   * The Alluxio Open Foundation licenses this work under the Apache License, version 2.0
   * (the "License"). You may not use this work except in compliance with the License, which is
   * available at www.apache.org/licenses/LICENSE-2.0
   *
   * This software is distributed on an "AS IS" basis, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
   * either express or implied, as more fully set forth in the License.
   *
   * See the NOTICE file distributed with this work for information regarding copyright ownership.
   */
  
  package alluxio.grpc;
  
  import io.grpc.CallOptions;
  import io.grpc.MethodDescriptor;
  import io.grpc.ServerMethodDefinition;
  import io.grpc.ServerServiceDefinition;
  import io.grpc.ServiceDescriptor;
  import io.grpc.internal.CompositeReadableBuffer;
  import io.grpc.internal.ReadableBuffer;
  import io.netty.buffer.ByteBuf;
  import io.netty.buffer.CompositeByteBuf;
  import io.netty.buffer.PooledByteBufAllocator;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.io.InputStream;
  import java.io.OutputStream;
  import java.lang.reflect.Constructor;
  import java.lang.reflect.Field;
  import java.util.ArrayList;
  import java.util.List;
  import java.util.Map;
  import java.util.Queue;
  
  /**
   * Utilities for gRPC message serialization.
   */
  public class GrpcSerializationUtils {
    public static final CallOptions.Key<MethodDescriptor> OVERRIDDEN_METHOD_DESCRIPTOR =
        CallOptions.Key.create("overridden method descriptor");
  
    private static final Logger LOG = LoggerFactory.getLogger(GrpcSerializationUtils.class);
  
    private static final int TAG_TYPE_BITS = 3;
  
    private static final String BUFFER_INPUT_STREAM_CLASS_NAME =
        "io.grpc.internal.ReadableBuffers$BufferInputStream";
    private static final String BUFFER_FIELD_NAME = "buffer";
    private static final String READABLE_BUFFERS_FIELD_NAME = "readableBuffers";
    private static final String NETTY_WRITABLE_BUFFER_CLASS_NAME =
        "io.grpc.netty.NettyWritableBuffer";
    private static final String NETTY_READABLE_BUFFER_CLASS_NAME =
        "io.grpc.netty.NettyReadableBuffer";
    private static final String BUFFER_CHAIN_OUTPUT_STREAM_CLASS_NAME =
        "io.grpc.internal.MessageFramer$BufferChainOutputStream";
    private static final String BUFFER_LIST_FIELD_NAME = "bufferList";
    private static final String CURRENT_FIELD_NAME = "current";
  
    private static Constructor<?> sNettyWritableBufferConstructor;
    private static Field sBufferList;
    private static Field sCompositeBuffers = null;
    private static Field sCurrent;
    private static Field sReadableBufferField = null;
    private static Field sReadableByteBuf = null;
    private static boolean sZeroCopySendSupported = true;
    private static boolean sZeroCopyReceiveSupported = true;
  
    static {
      try {
        sReadableBufferField = getPrivateField(BUFFER_INPUT_STREAM_CLASS_NAME, BUFFER_FIELD_NAME);
      } catch (Exception e) {
        LOG.warn("Cannot get gRPC input stream buffer, zero copy send will be disabled.", e);
        sZeroCopySendSupported = false;
      }
      try {
        sNettyWritableBufferConstructor =
            getPrivateConstructor(NETTY_WRITABLE_BUFFER_CLASS_NAME, ByteBuf.class);
        sBufferList = getPrivateField(BUFFER_CHAIN_OUTPUT_STREAM_CLASS_NAME, BUFFER_LIST_FIELD_NAME);
        sCurrent = getPrivateField(BUFFER_CHAIN_OUTPUT_STREAM_CLASS_NAME, CURRENT_FIELD_NAME);
        sCompositeBuffers =
            getPrivateField(CompositeReadableBuffer.class.getName(), READABLE_BUFFERS_FIELD_NAME);
        sReadableByteBuf = getPrivateField(NETTY_READABLE_BUFFER_CLASS_NAME, BUFFER_FIELD_NAME);
      } catch (Exception e) {
        LOG.warn("Cannot get gRPC output stream buffer, zero copy receive will be disabled.", e);
        sZeroCopyReceiveSupported = false;
      }
    }
  
    private static Field getPrivateField(String className, String fieldName)
        throws NoSuchFieldException, ClassNotFoundException {
      Class<?> declaringClass = Class.forName(className);
      Field field = declaringClass.getDeclaredField(fieldName);
      field.setAccessible(true);
      return field;
    }
  
    private static Constructor<?> getPrivateConstructor(String className, Class<?> ... parameterTypes)
        throws ClassNotFoundException, NoSuchMethodException {
      Class<?> declaringClass = Class.forName(className);
      Constructor<?> constructor = declaringClass.getDeclaredConstructor(parameterTypes);
      constructor.setAccessible(true);
      return constructor;
    }
  
    /**
     * Makes a gRPC tag for a field.
     *
     * @param fieldNumber field number
     * @param wireType wire type of the field
     * @return the gRPC tag
     */
    public static int makeTag(final int fieldNumber, final int wireType) {
      // This is a public version of WireFormat.makeTag.
      return (fieldNumber << TAG_TYPE_BITS) | wireType;
    }
  
    /**
     * Gets a buffer directly from a gRPC input stream.
     *
     * @param stream the input stream
     * @return the raw data buffer
     */
    public static ReadableBuffer getBufferFromStream(InputStream stream) {
      if (!sZeroCopyReceiveSupported
          || !stream.getClass().equals(sReadableBufferField.getDeclaringClass())) {
        return null;
      }
      try {
        return (ReadableBuffer) sReadableBufferField.get(stream);
      } catch (Exception e) {
        LOG.warn("Failed to get data buffer from stream.", e);
        return null;
      }
    }
  
    /**
     * Gets a Netty buffer directly from a gRPC ReadableBuffer.
     *
     * @param buffer the input buffer
     * @return the raw ByteBuf, or null if the ByteBuf cannot be extracted
     */
    public static ByteBuf getByteBufFromReadableBuffer(ReadableBuffer buffer) {
      if (!sZeroCopyReceiveSupported) {
        return null;
      }
      try {
        if (buffer instanceof CompositeReadableBuffer) {
          // TODO(elega) grpc introduced native protobuf zero copy since 1.39.0
          //  https://github.com/grpc/grpc-java/pull/8102/files
          //  replace the following with
          //    return Unpooled.wrappedBuffer(buffer.getByteBuffer());
          Queue<ReadableBuffer> buffers = (Queue<ReadableBuffer>) sCompositeBuffers.get(buffer);
          if (buffers.size() == 1) {
            return getByteBufFromReadableBuffer(buffers.peek());
          } else {
            CompositeByteBuf buf = PooledByteBufAllocator.DEFAULT.compositeBuffer();
            for (ReadableBuffer readableBuffer : buffers) {
              ByteBuf subBuffer = getByteBufFromReadableBuffer(readableBuffer);
              if (subBuffer == null) {
                return null;
              }
              buf.addComponent(true, subBuffer);
            }
            return buf;
          }
        } else if (buffer.getClass().equals(sReadableByteBuf.getDeclaringClass())) {
          return (ByteBuf) sReadableByteBuf.get(buffer);
        }
      } catch (Exception e) {
        LOG.warn("Failed to get data buffer from stream: {}.", e.toString());
        return null;
      }
      return null;
    }
  
    /**
     * Add the given buffers directly to the gRPC output stream.
     *
     * @param buffers the buffers to be added
     * @param stream the output stream
     * @return whether the buffers are added successfully
     */
    public static boolean addBuffersToStream(ByteBuf[] buffers, OutputStream stream) {
      if (!sZeroCopySendSupported || !stream.getClass().equals(sBufferList.getDeclaringClass())) {
        return false;
      }
      try {
        if (sCurrent.get(stream) != null) {
          return false;
        }
        for (ByteBuf buffer : buffers) {
          Object nettyBuffer = sNettyWritableBufferConstructor.newInstance(buffer);
          List list = (List) sBufferList.get(stream);
          list.add(nettyBuffer);
          buffer.retain();
          sCurrent.set(stream, nettyBuffer);
        }
        return true;
      } catch (Exception e) {
        LOG.warn("Failed to add data buffer to stream: {}.", e.toString());
        return false;
      }
    }
  
    /**
     * Creates a service definition that uses custom marshallers.
     *
     * @param service the service to intercept
     * @param marshallers a map that specifies which marshaller to use for each method
     * @return the new service definition
     */
    public static ServerServiceDefinition overrideMethods(
        final ServerServiceDefinition service,
        final Map<MethodDescriptor, MethodDescriptor> marshallers) {
      List<ServerMethodDefinition<?, ?>> newMethods = new ArrayList<ServerMethodDefinition<?, ?>>();
      List<MethodDescriptor<?, ?>> newDescriptors = new ArrayList<MethodDescriptor<?, ?>>();
      // intercepts the descriptors
      for (final ServerMethodDefinition<?, ?> definition : service.getMethods()) {
        ServerMethodDefinition<?, ?> newMethod = interceptMethod(definition, marshallers);
        newDescriptors.add(newMethod.getMethodDescriptor());
        newMethods.add(newMethod);
      }
      // builds the new service descriptor
      final ServerServiceDefinition.Builder serviceBuilder = ServerServiceDefinition
          .builder(new ServiceDescriptor(service.getServiceDescriptor().getName(), newDescriptors));
      // creates the new service definition
      for (ServerMethodDefinition<?, ?> definition : newMethods) {
        serviceBuilder.addMethod(definition);
      }
      return serviceBuilder.build();
    }
  
    private static <ReqT, RespT>  ServerMethodDefinition<ReqT, RespT> interceptMethod(
        final ServerMethodDefinition<ReqT, RespT> definition,
        final Map<MethodDescriptor, MethodDescriptor> newMethods) {
      MethodDescriptor<ReqT, RespT> descriptor = definition.getMethodDescriptor();
      MethodDescriptor newMethod = newMethods.get(descriptor);
      if (newMethod != null) {
        return ServerMethodDefinition.create(newMethod, definition.getServerCallHandler());
      }
      return definition;
    }
  }
  