Merge pull request #493 from VelocityPowered/feature/combine-varint-prefix-with-compression

Further packet handling performance improvements

proxy/src/main/java/com/velocitypowered/proxy/connection/MinecraftConnection.java

@@ -36,12 +36,13 @@ import com.velocitypowered.proxy.VelocityServer;
 import com.velocitypowered.proxy.connection.client.HandshakeSessionHandler;
 import com.velocitypowered.proxy.connection.client.LoginSessionHandler;
 import com.velocitypowered.proxy.connection.client.StatusSessionHandler;
+import com.velocitypowered.proxy.network.Connections;
 import com.velocitypowered.proxy.protocol.MinecraftPacket;
 import com.velocitypowered.proxy.protocol.StateRegistry;
 import com.velocitypowered.proxy.protocol.netty.MinecraftCipherDecoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftCipherEncoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftCompressDecoder;
-import com.velocitypowered.proxy.protocol.netty.MinecraftCompressEncoder;
+import com.velocitypowered.proxy.protocol.netty.MinecraftCompressorAndLengthEncoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftDecoder;
 import com.velocitypowered.proxy.protocol.netty.MinecraftEncoder;
 import com.velocitypowered.proxy.util.except.QuietDecoderException;
@@ -402,8 +403,8 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
     } else {
       MinecraftCompressDecoder decoder = (MinecraftCompressDecoder) channel.pipeline()
           .get(COMPRESSION_DECODER);
-      MinecraftCompressEncoder encoder = (MinecraftCompressEncoder) channel.pipeline()
-          .get(COMPRESSION_ENCODER);
+      MinecraftCompressorAndLengthEncoder encoder =
+          (MinecraftCompressorAndLengthEncoder) channel.pipeline().get(COMPRESSION_ENCODER);
       if (decoder != null && encoder != null) {
         decoder.setThreshold(threshold);
         encoder.setThreshold(threshold);
@@ -411,9 +412,10 @@ public class MinecraftConnection extends ChannelInboundHandlerAdapter {
         int level = server.getConfiguration().getCompressionLevel();
         VelocityCompressor compressor = Natives.compress.get().create(level);
 
-        encoder = new MinecraftCompressEncoder(threshold, compressor);
+        encoder = new MinecraftCompressorAndLengthEncoder(threshold, compressor);
         decoder = new MinecraftCompressDecoder(threshold, compressor);
 
+        channel.pipeline().remove(FRAME_ENCODER);
         channel.pipeline().addBefore(MINECRAFT_DECODER, COMPRESSION_DECODER, decoder);
         channel.pipeline().addBefore(MINECRAFT_ENCODER, COMPRESSION_ENCODER, encoder);
       }

proxy/src/main/java/com/velocitypowered/proxy/protocol/ProtocolUtils.java

@@ -120,7 +120,20 @@ public enum ProtocolUtils {
    * @param value the integer to write
    */
   public static void writeVarInt(ByteBuf buf, int value) {
-    // Optimization: focus on 1-3 byte VarInts as they are the most common
+    // Peel the one and two byte count cases explicitly as they are the most common VarInt sizes
+    // that the proxy will write, to improve inlining.
+    if ((value & (0xFFFFFFFF << 7)) == 0) {
+      buf.writeByte(value);
+    } else if ((value & (0xFFFFFFFF << 14)) == 0) {
+      int w = (value & 0x7F | 0x80) << 8 | (value >>> 7);
+      buf.writeShort(w);
+    } else {
+      writeVarIntFull(buf, value);
+    }
+  }
+
+  private static void writeVarIntFull(ByteBuf buf, int value) {
+    // See https://steinborn.me/posts/performance/how-fast-can-you-write-a-varint/
     if ((value & (0xFFFFFFFF << 7)) == 0) {
       buf.writeByte(value);
     } else if ((value & (0xFFFFFFFF << 14)) == 0) {
@@ -129,22 +142,28 @@ public enum ProtocolUtils {
     } else if ((value & (0xFFFFFFFF << 21)) == 0) {
       int w = (value & 0x7F | 0x80) << 16 | ((value >>> 7) & 0x7F | 0x80) << 8 | (value >>> 14);
       buf.writeMedium(w);
+    } else if ((value & (0xFFFFFFFF << 28)) == 0) {
+      int w = (value & 0x7F | 0x80) << 24 | (((value >>> 7) & 0x7F | 0x80) << 16)
+          | ((value >>> 14) & 0x7F | 0x80) << 8 | (value >>> 21);
+      buf.writeInt(w);
     } else {
-      // 4 and 5 byte VarInts aren't common so split those cases off
-      writeVarIntUncommon(buf, value);
+      int w = (value & 0x7F | 0x80) << 24 | ((value >>> 7) & 0x7F | 0x80) << 16
+          | ((value >>> 14) & 0x7F | 0x80) << 8 | ((value >>> 21) & 0x7F | 0x80);
+      buf.writeInt(w);
+      buf.writeByte(value >>> 28);
     }
   }
 
-  private static void writeVarIntUncommon(ByteBuf buf, int value) {
-    while (true) {
-      if ((value & 0xFFFFFF80) == 0) {
-        buf.writeByte(value);
-        return;
-      }
-
-      buf.writeByte(value & 0x7F | 0x80);
-      value >>>= 7;
-    }
+  /**
+   * Writes the specified {@code value} as a 21-bit Minecraft VarInt to the specified {@code buf}.
+   * The upper 11 bits will be discarded.
+   * @param buf the buffer to read from
+   * @param value the integer to write
+   */
+  public static void write21BitVarInt(ByteBuf buf, int value) {
+    // See https://steinborn.me/posts/performance/how-fast-can-you-write-a-varint/
+    int w = (value & 0x7F | 0x80) << 16 | ((value >>> 7) & 0x7F | 0x80) << 8 | (value >>> 14);
+    buf.writeMedium(w);
   }
 
   public static String readString(ByteBuf buf) {

proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressEncoder.java

@@ -1,78 +0,0 @@
-/*
- * Copyright (C) 2018 Velocity Contributors
- *
- * This program is free software: you can redistribute it and/or modify
- * it under the terms of the GNU General Public License as published by
- * the Free Software Foundation, either version 3 of the License, or
- * (at your option) any later version.
- *
- * This program is distributed in the hope that it will be useful,
- * but WITHOUT ANY WARRANTY; without even the implied warranty of
- * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
- * GNU General Public License for more details.
- *
- * You should have received a copy of the GNU General Public License
- * along with this program.  If not, see <https://www.gnu.org/licenses/>.
- */
-
-package com.velocitypowered.proxy.protocol.netty;
-
-import com.velocitypowered.natives.compression.VelocityCompressor;
-import com.velocitypowered.natives.util.MoreByteBufUtils;
-import com.velocitypowered.proxy.protocol.ProtocolUtils;
-import io.netty.buffer.ByteBuf;
-import io.netty.channel.ChannelHandlerContext;
-import io.netty.handler.codec.MessageToByteEncoder;
-
-public class MinecraftCompressEncoder extends MessageToByteEncoder<ByteBuf> {
-
-  private int threshold;
-  private final VelocityCompressor compressor;
-
-  public MinecraftCompressEncoder(int threshold, VelocityCompressor compressor) {
-    this.threshold = threshold;
-    this.compressor = compressor;
-  }
-
-  @Override
-  protected void encode(ChannelHandlerContext ctx, ByteBuf msg, ByteBuf out) throws Exception {
-    int uncompressed = msg.readableBytes();
-    if (uncompressed < threshold) {
-      // Under the threshold, there is nothing to do.
-      ProtocolUtils.writeVarInt(out, 0);
-      out.writeBytes(msg);
-    } else {
-      ProtocolUtils.writeVarInt(out, uncompressed);
-      ByteBuf compatibleIn = MoreByteBufUtils.ensureCompatible(ctx.alloc(), compressor, msg);
-      try {
-        compressor.deflate(compatibleIn, out);
-      } finally {
-        compatibleIn.release();
-      }
-    }
-  }
-
-  @Override
-  protected ByteBuf allocateBuffer(ChannelHandlerContext ctx, ByteBuf msg, boolean preferDirect)
-      throws Exception {
-    // We allocate bytes to be compressed plus 1 byte. This covers two cases:
-    //
-    // - Compression
-    //    According to https://github.com/ebiggers/libdeflate/blob/master/libdeflate.h#L103,
-    //    if the data compresses well (and we do not have some pathological case) then the maximum
-    //    size the compressed size will ever be is the input size minus one.
-    // - Uncompressed
-    //    This is fairly obvious - we will then have one more than the uncompressed size.
-    int initialBufferSize = msg.readableBytes() + 1;
-    return MoreByteBufUtils.preferredBuffer(ctx.alloc(), compressor, initialBufferSize);
-  }
-
-  @Override
-  public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
-    compressor.close();
-  }
-
-  public void setThreshold(int threshold) {
-    this.threshold = threshold;
-  }
-}

proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftCompressorAndLengthEncoder.java

@@ -0,0 +1,133 @@
+/*
+ * Copyright (C) 2018 Velocity Contributors
+ *
+ * This program is free software: you can redistribute it and/or modify
+ * it under the terms of the GNU General Public License as published by
+ * the Free Software Foundation, either version 3 of the License, or
+ * (at your option) any later version.
+ *
+ * This program is distributed in the hope that it will be useful,
+ * but WITHOUT ANY WARRANTY; without even the implied warranty of
+ * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+ * GNU General Public License for more details.
+ *
+ * You should have received a copy of the GNU General Public License
+ * along with this program.  If not, see <https://www.gnu.org/licenses/>.
+ */
+
+package com.velocitypowered.proxy.protocol.netty;
+
+import static com.velocitypowered.proxy.protocol.netty.MinecraftVarintLengthEncoder.IS_JAVA_CIPHER;
+
+import com.velocitypowered.natives.compression.VelocityCompressor;
+import com.velocitypowered.natives.util.MoreByteBufUtils;
+import com.velocitypowered.proxy.protocol.ProtocolUtils;
+import io.netty.buffer.ByteBuf;
+import io.netty.channel.ChannelHandlerContext;
+import io.netty.handler.codec.MessageToByteEncoder;
+import java.util.zip.DataFormatException;
+
+public class MinecraftCompressorAndLengthEncoder extends MessageToByteEncoder<ByteBuf> {
+
+  private static final boolean MUST_USE_SAFE_AND_SLOW_COMPRESSION_HANDLING =
+      Boolean.getBoolean("velocity.increased-compression-cap");
+
+  private int threshold;
+  private final VelocityCompressor compressor;
+
+  public MinecraftCompressorAndLengthEncoder(int threshold, VelocityCompressor compressor) {
+    this.threshold = threshold;
+    this.compressor = compressor;
+  }
+
+  @Override
+  protected void encode(ChannelHandlerContext ctx, ByteBuf msg, ByteBuf out) throws Exception {
+    int uncompressed = msg.readableBytes();
+    if (uncompressed < threshold) {
+      // Under the threshold, there is nothing to do.
+      ProtocolUtils.writeVarInt(out, uncompressed + 1);
+      ProtocolUtils.writeVarInt(out, 0);
+      out.writeBytes(msg);
+    } else {
+      if (MUST_USE_SAFE_AND_SLOW_COMPRESSION_HANDLING) {
+        handleCompressedSafe(ctx, msg, out);
+      } else {
+        handleCompressedFast(ctx, msg, out);
+      }
+    }
+  }
+
+  private void handleCompressedFast(ChannelHandlerContext ctx, ByteBuf msg, ByteBuf out)
+      throws DataFormatException {
+    int uncompressed = msg.readableBytes();
+
+    ProtocolUtils.write21BitVarInt(out, 0); // Dummy packet length
+    ProtocolUtils.writeVarInt(out, uncompressed);
+    ByteBuf compatibleIn = MoreByteBufUtils.ensureCompatible(ctx.alloc(), compressor, msg);
+
+    int startCompressed = out.writerIndex();
+    try {
+      compressor.deflate(compatibleIn, out);
+    } finally {
+      compatibleIn.release();
+    }
+    int compressedLength = out.writerIndex() - startCompressed;
+    if (compressedLength >= 1 << 21) {
+      throw new DataFormatException("The server sent a very large (over 2MiB compressed) packet. "
+          + "Please restart Velocity with the JVM flag -Dvelocity.increased-compression-cap=true "
+          + "to fix this issue.");
+    }
+
+    int writerIndex = out.writerIndex();
+    int packetLength = out.readableBytes() - 3;
+    out.writerIndex(0);
+    ProtocolUtils.write21BitVarInt(out, packetLength); // Rewrite packet length
+    out.writerIndex(writerIndex);
+  }
+
+  private void handleCompressedSafe(ChannelHandlerContext ctx, ByteBuf msg, ByteBuf out)
+      throws DataFormatException {
+    int uncompressed = msg.readableBytes();
+    ByteBuf tmpBuf = MoreByteBufUtils.preferredBuffer(ctx.alloc(), compressor, uncompressed - 1);
+    try {
+      ProtocolUtils.writeVarInt(tmpBuf, uncompressed);
+      ByteBuf compatibleIn = MoreByteBufUtils.ensureCompatible(ctx.alloc(), compressor, msg);
+      try {
+        compressor.deflate(compatibleIn, tmpBuf);
+      } finally {
+        compatibleIn.release();
+      }
+
+      ProtocolUtils.writeVarInt(out, tmpBuf.readableBytes());
+      out.writeBytes(tmpBuf);
+    } finally {
+      tmpBuf.release();
+    }
+  }
+
+  @Override
+  protected ByteBuf allocateBuffer(ChannelHandlerContext ctx, ByteBuf msg, boolean preferDirect)
+      throws Exception {
+    int uncompressed = msg.readableBytes();
+    if (uncompressed < threshold) {
+      int finalBufferSize = uncompressed + 1;
+      finalBufferSize += ProtocolUtils.varIntBytes(finalBufferSize);
+      return IS_JAVA_CIPHER
+          ? ctx.alloc().heapBuffer(finalBufferSize)
+          : ctx.alloc().directBuffer(finalBufferSize);
+    }
+
+    // (maximum data length after compression) + packet length varint + uncompressed data varint
+    int initialBufferSize = (uncompressed - 1) + 3 + ProtocolUtils.varIntBytes(uncompressed);
+    return MoreByteBufUtils.preferredBuffer(ctx.alloc(), compressor, initialBufferSize);
+  }
+
+  @Override
+  public void handlerRemoved(ChannelHandlerContext ctx) throws Exception {
+    compressor.close();
+  }
+
+  public void setThreshold(int threshold) {
+    this.threshold = threshold;
+  }
+}

proxy/src/main/java/com/velocitypowered/proxy/protocol/netty/MinecraftVarintLengthEncoder.java

@@ -29,7 +29,7 @@ import io.netty.handler.codec.MessageToByteEncoder;
 public class MinecraftVarintLengthEncoder extends MessageToByteEncoder<ByteBuf> {
 
   public static final MinecraftVarintLengthEncoder INSTANCE = new MinecraftVarintLengthEncoder();
-  private static final boolean IS_JAVA_CIPHER = Natives.cipher.get() == JavaVelocityCipher.FACTORY;
+  public static final boolean IS_JAVA_CIPHER = Natives.cipher.get() == JavaVelocityCipher.FACTORY;
 
   private MinecraftVarintLengthEncoder() {
   }

proxy/src/test/java/com/velocitypowered/proxy/protocol/ProtocolUtilsTest.java

@@ -71,6 +71,20 @@ public class ProtocolUtilsTest {
     assertEquals(test, ProtocolUtils.readVarIntSafely(buf));
   }
 
+  @Test
+  void test3Bytes() {
+    ByteBuf buf = Unpooled.buffer(5);
+    for (int i = 0; i < 2097152; i += 31) {
+      writeReadTest3Bytes(buf, i);
+    }
+  }
+
+  private void writeReadTest3Bytes(ByteBuf buf, int test) {
+    buf.clear();
+    ProtocolUtils.write21BitVarInt(buf, test);
+    assertEquals(test, ProtocolUtils.readVarInt(buf));
+  }
+
   @Test
   void testBytesWrittenAtBitBoundaries() {
     ByteBuf varintNew = Unpooled.buffer(5);