Add compressed payload pass-through mode for SR Proxy

third-party/thrift/src/thrift/lib/cpp2/async/RocketClientChannelBase.cpp

@@ -45,6 +45,7 @@
 #include <thrift/lib/cpp2/transport/rocket/FdSocket.h>
 #include <thrift/lib/cpp2/transport/rocket/RocketException.h>
 #include <thrift/lib/cpp2/transport/rocket/client/RocketClient.h>
+#include <thrift/lib/cpp2/transport/rocket/compression/CompressionManager.h>
 #include <thrift/lib/cpp2/transport/rocket/payload/PayloadSerializer.h>
 #include <thrift/lib/thrift/gen-cpp2/RpcMetadata_constants.h>
 #include <thrift/lib/thrift/gen-cpp2/RpcMetadata_types.h>
@@ -978,6 +979,55 @@ void RocketClientChannelBase::sendThriftRequest(
       std::move(frameworkMetadata),
       hasCustomCompressor());
 
+  // If the payload is already compressed (pre-compressed pass-through mode),
+  // check whether we can forward the compressed buffer as-is. This requires:
+  //   1. The outbound channel either has no compression configured, or its
+  //      configured algorithm matches the inbound pre-compressed algorithm.
+  //   2. The target server supports the inbound algorithm.
+  // If both conditions are met, set metadata.compression to the inbound
+  // algorithm (so the receiver knows to decompress) and skip re-compression.
+  // Otherwise, decompress the buffer and let normal compression flow proceed.
+  bool skipCompression = false;
+  auto preCompressedAlgo = rpcOptions.getPreCompressedAlgorithm();
+  if (preCompressedAlgo != CompressionAlgorithm::NONE) {
+    // Check if the outbound channel already has compression configured
+    // (via makeRequestRpcMetadata from THeader's DesiredCompressionConfig).
+    auto outboundCompression = metadata.compression();
+    bool outboundAlgoConflicts = outboundCompression.has_value() &&
+        *outboundCompression != CompressionAlgorithm::NONE &&
+        *outboundCompression != preCompressedAlgo;
+
+    if (outboundAlgoConflicts) {
+      // Outbound wants a different algorithm than what we have --
+      // decompress so packWithFds can re-compress with the outbound algo.
+      buf = rocket::CompressionManager().uncompressBuffer(
+          std::move(buf), preCompressedAlgo);
+    } else {
+      // No conflicting outbound algo. Check if target supports inbound algo.
+      bool targetSupportsAlgo = false;
+      switch (preCompressedAlgo) {
+        case CompressionAlgorithm::ZSTD:
+        case CompressionAlgorithm::ZSTD_LESS:
+        case CompressionAlgorithm::ZSTD_MORE:
+          targetSupportsAlgo = getRocketClientImpl().getServerZstdSupported();
+          break;
+        default:
+          // No reliable detection for non-ZSTD algorithms yet
+          targetSupportsAlgo = false;
+          break;
+      }
+      if (targetSupportsAlgo) {
+        metadata.compression() = preCompressedAlgo;
+        skipCompression = true;
+      } else {
+        // Target doesn't support this algorithm -- decompress and let
+        // normal compression logic in packWithFds handle it.
+        buf = rocket::CompressionManager().uncompressBuffer(
+            std::move(buf), preCompressedAlgo);
+      }
+    }
+  }
+
   size_t requestSerializedSize;
   // Avoid unnecessary computation for streaming methods.
   if constexpr (std::is_same_v<Callback, RequestClientCallback::Ptr>) {
@@ -991,7 +1041,8 @@ void RocketClientChannelBase::sendThriftRequest(
         rpcOptions.copySocketFdsToSend(),
         encodeMetadataUsingBinary(),
         getRocketClientImpl().getTransportWrapper(),
-        getRocketClientImpl().getIOBufFactory());
+        getRocketClientImpl().getIOBufFactory(),
+        skipCompression);
     requestPayload.setDataFirstFieldAlignment(
         rpcOptions.getFrameRelativeDataAlignment());
     if (metadata.protocol()) {

third-party/thrift/src/thrift/lib/cpp2/async/RpcOptions.cpp

@@ -305,6 +305,15 @@ const std::shared_ptr<void>& RpcOptions::getMetricsToCollect() const {
   return metricsToCollect_;
 }
 
+RpcOptions& RpcOptions::setPreCompressedAlgorithm(CompressionAlgorithm algo) {
+  preCompressedAlgorithm_ = algo;
+  return *this;
+}
+
+CompressionAlgorithm RpcOptions::getPreCompressedAlgorithm() const {
+  return preCompressedAlgorithm_;
+}
+
 RpcOptions& RpcOptions::setRoutingObjectiveKey(
     std::string routingObjectiveKey) {
   routingObjectiveKey_ = std::move(routingObjectiveKey);

third-party/thrift/src/thrift/lib/cpp2/async/RpcOptions.h

@@ -183,6 +183,19 @@ class RpcOptions {
   RpcOptions& setForceSyncOnFiber(bool forceSyncOnFiber);
   bool getForceSyncOnFiber() const;
 
+  /**
+   * Indicates that the payload buffer in SerializedRequest is already
+   * compressed with the specified algorithm. When set, the channel layer
+   * will set metadata.compression for the wire but skip re-compressing
+   * the buffer in packWithFds().
+   *
+   * This is used by proxies (e.g., SR Proxy) to forward compressed payloads
+   * without decompression/recompression when the target supports the same
+   * compression algorithm.
+   */
+  RpcOptions& setPreCompressedAlgorithm(CompressionAlgorithm algo);
+  CompressionAlgorithm getPreCompressedAlgorithm() const;
+
   /**
    * Set routing objective key for global routing optimization.
    *
@@ -258,6 +271,9 @@ class RpcOptions {
 
   bool forceSyncOnFiber_{false};
 
+  // When set, indicates the payload is already compressed with this algorithm.
+  CompressionAlgorithm preCompressedAlgorithm_{CompressionAlgorithm::NONE};
+
   // Metrics to be sent back to the client
   std::shared_ptr<void> metricsToCollect_;
 

third-party/thrift/src/thrift/lib/cpp2/transport/rocket/payload/ChecksumPayloadSerializerStrategy.h

@@ -100,7 +100,8 @@ class ChecksumPayloadSerializerStrategy final
       folly::SocketFds fds,
       bool encodeMetadataUsingBinary,
       folly::AsyncTransport* transport,
-      folly::IOBufFactory* ioBufFactory = nullptr) {
+      folly::IOBufFactory* ioBufFactory = nullptr,
+      bool skipCompression = false) {
     if (payload != nullptr) {
       if (auto checksumOpt =
               calculateChecksum(*payload, metadata->checksum())) {
@@ -113,7 +114,8 @@ class ChecksumPayloadSerializerStrategy final
         std::move(fds),
         encodeMetadataUsingBinary,
         transport,
-        ioBufFactory);
+        ioBufFactory,
+        skipCompression);
   }
 
   template <class PayloadType>

third-party/thrift/src/thrift/lib/cpp2/transport/rocket/payload/CustomCompressionPayloadSerializerStrategy.h

@@ -102,10 +102,13 @@ class CustomCompressionPayloadSerializerStrategy final
       folly::SocketFds fds,
       bool encodeMetadataUsingBinary,
       folly::AsyncTransport* transport,
-      folly::IOBufFactory* ioBufFactory = nullptr) {
-    if (auto compress = metadata->compression()) {
-      if (*compress == CompressionAlgorithm::CUSTOM) {
-        payload = customCompressBuffer(std::move(payload));
+      folly::IOBufFactory* ioBufFactory = nullptr,
+      bool skipCompression = false) {
+    if (!skipCompression) {
+      if (auto compress = metadata->compression()) {
+        if (*compress == CompressionAlgorithm::CUSTOM) {
+          payload = customCompressBuffer(std::move(payload));
+        }
       }
     }
 
@@ -115,7 +118,8 @@ class CustomCompressionPayloadSerializerStrategy final
         std::move(fds),
         encodeMetadataUsingBinary,
         transport,
-        ioBufFactory);
+        ioBufFactory,
+        skipCompression);
   }
 
   template <class PayloadType>

third-party/thrift/src/thrift/lib/cpp2/transport/rocket/payload/DefaultPayloadSerializerStrategy.cpp

@@ -54,14 +54,17 @@ rocket::Payload DefaultPayloadSerializerStrategy::packWithFds(
     folly::SocketFds fds,
     bool encodeMetadataUsingBinary,
     folly::AsyncTransport* transport,
-    folly::IOBufFactory* ioBufFactory) {
-  if (auto compression = metadata->compression_ref()) {
-    const auto compressionAlgorithm = *compression;
-    if (compressionAlgorithm != CompressionAlgorithm::NONE &&
-        compressionAlgorithm != CompressionAlgorithm::CUSTOM) {
-      // Custom compression is supported in
-      // CustomCompressionPayloadSerializerStrategy
-      payload = compressBuffer(std::move(payload), compressionAlgorithm);
+    folly::IOBufFactory* ioBufFactory,
+    bool skipCompression) {
+  if (!skipCompression) {
+    if (auto compression = metadata->compression_ref()) {
+      const auto compressionAlgorithm = *compression;
+      if (compressionAlgorithm != CompressionAlgorithm::NONE &&
+          compressionAlgorithm != CompressionAlgorithm::CUSTOM) {
+        // Custom compression is supported in
+        // CustomCompressionPayloadSerializerStrategy
+        payload = compressBuffer(std::move(payload), compressionAlgorithm);
+      }
     }
   }
 
@@ -85,7 +88,8 @@ DefaultPayloadSerializerStrategy::packWithFds<RequestRpcMetadata>(
     folly::SocketFds,
     bool,
     folly::AsyncTransport*,
-    folly::IOBufFactory*);
+    folly::IOBufFactory*,
+    bool);
 
 template rocket::Payload
 DefaultPayloadSerializerStrategy::packWithFds<ResponseRpcMetadata>(
@@ -94,7 +98,8 @@ DefaultPayloadSerializerStrategy::packWithFds<ResponseRpcMetadata>(
     folly::SocketFds,
     bool,
     folly::AsyncTransport*,
-    folly::IOBufFactory*);
+    folly::IOBufFactory*,
+    bool);
 
 template rocket::Payload
 DefaultPayloadSerializerStrategy::packWithFds<StreamPayloadMetadata>(
@@ -103,7 +108,8 @@ DefaultPayloadSerializerStrategy::packWithFds<StreamPayloadMetadata>(
     folly::SocketFds,
     bool,
     folly::AsyncTransport*,
-    folly::IOBufFactory*);
+    folly::IOBufFactory*,
+    bool);
 
 bool DefaultPayloadSerializerStrategy::
     canSerializeMetadataIntoDataBufferHeadroom(

third-party/thrift/src/thrift/lib/cpp2/transport/rocket/payload/DefaultPayloadSerializerStrategy.h

@@ -67,7 +67,8 @@ class DefaultPayloadSerializerStrategy final
       folly::SocketFds fds,
       bool encodeMetadataUsingBinary,
       folly::AsyncTransport* transport,
-      folly::IOBufFactory* ioBufFactory = nullptr);
+      folly::IOBufFactory* ioBufFactory = nullptr,
+      bool skipCompression = false);
 
   template <typename T>
   std::unique_ptr<folly::IOBuf> packCompact(const T& data) {

third-party/thrift/src/thrift/lib/cpp2/transport/rocket/payload/PayloadSerializer.h

@@ -240,7 +240,8 @@ class PayloadSerializer : public folly::hazptr_obj_base<PayloadSerializer> {
       folly::SocketFds fds,
       bool encodeMetadataUsingBinary,
       folly::AsyncTransport* transport,
-      folly::IOBufFactory* ioBufFactory = nullptr) {
+      folly::IOBufFactory* ioBufFactory = nullptr,
+      bool skipCompression = false) {
     if (!supportsChecksum()) {
       if (metadata->checksum().has_value() &&
           metadata->checksum()->algorithm().value() !=
@@ -259,7 +260,8 @@ class PayloadSerializer : public folly::hazptr_obj_base<PayloadSerializer> {
           std::move(fds),
           encodeMetadataUsingBinary,
           transport,
-          ioBufFactory);
+          ioBufFactory,
+          skipCompression);
     });
   }
 

third-party/thrift/src/thrift/lib/cpp2/transport/rocket/payload/test/PayloadSerializerTest.cpp

@@ -116,6 +116,82 @@ TEST(PayloadSerializerTest, TestMakeCustomCompression) {
   testPackAndUnpackWithCompactProtocol(ps);
 }
 
+TEST(PayloadSerializerTest, TestSkipCompressionPreservesBuffer) {
+  if (!folly::kIsLinux) {
+    return;
+  }
+
+  // Test that packWithFds with skipCompression=true preserves the buffer
+  // unchanged while still writing the compression algorithm into metadata.
+  // This is used by the SR Proxy pass-through path: the payload is already
+  // compressed, so we set metadata.compression for the receiver but skip
+  // re-compressing the buffer.
+
+  PayloadSerializer::reset();
+  PayloadSerializer::initialize(DefaultPayloadSerializerStrategy());
+  auto& serializer = *PayloadSerializer::getInstance().get();
+
+  std::string const plaintext = "hello world - test payload for compression";
+
+  // First, compress the buffer to simulate a pre-compressed payload
+  auto preCompressed = serializer.compressBuffer(
+      folly::IOBuf::fromString(plaintext), CompressionAlgorithm::ZSTD);
+  auto preCompressedCopy = preCompressed->clone();
+
+  // Pack with skipCompression=true: metadata says ZSTD, but buffer is NOT
+  // re-compressed
+  RequestRpcMetadata metadata;
+  metadata.protocol() = ProtocolId::COMPACT;
+  metadata.compression() = CompressionAlgorithm::ZSTD;
+  auto payload = serializer.packWithFds(
+      &metadata,
+      std::move(preCompressed),
+      folly::SocketFds(),
+      false, /* encodeMetadataUsingBinary */
+      nullptr, /* transport */
+      nullptr, /* ioBufFactory */
+      true /* skipCompression */);
+
+  // Unpack the payload -- the receiver sees metadata.compression=ZSTD and
+  // will decompress the buffer
+  auto unpacked = serializer.unpack<RequestPayload>(std::move(payload), false);
+  ASSERT_FALSE(unpacked.hasException());
+  auto& result = unpacked.value();
+
+  // The metadata should indicate ZSTD compression
+  EXPECT_TRUE(result.metadata.compression().has_value());
+  EXPECT_EQ(*result.metadata.compression(), CompressionAlgorithm::ZSTD);
+
+  // The unpacked payload should be the pre-compressed bytes (since the
+  // receiver's unpack sees compression=ZSTD and decompresses)
+  // So we should get back the original plaintext
+  EXPECT_EQ(result.payload->toString(), plaintext);
+
+  // Now do the same WITHOUT skipCompression (normal path) -- the buffer
+  // would be double-compressed and decompression would yield garbage or fail
+  auto preCompressed2 = serializer.compressBuffer(
+      folly::IOBuf::fromString(plaintext), CompressionAlgorithm::ZSTD);
+
+  RequestRpcMetadata metadata2;
+  metadata2.protocol() = ProtocolId::COMPACT;
+  metadata2.compression() = CompressionAlgorithm::ZSTD;
+  auto payload2 = serializer.packWithFds(
+      &metadata2,
+      std::move(preCompressed2),
+      folly::SocketFds(),
+      false,
+      nullptr,
+      nullptr,
+      false /* skipCompression = false, will double-compress */);
+
+  auto unpacked2 =
+      serializer.unpack<RequestPayload>(std::move(payload2), false);
+  ASSERT_FALSE(unpacked2.hasException());
+
+  // Double-compressed then single-decompressed: result should NOT be plaintext
+  EXPECT_NE(unpacked2.value().payload->toString(), plaintext);
+}
+
 TEST(PayloadSerializerTest, TestCompressionAndUncompression) {
   if (!folly::kIsLinux) {
     // on non-linux platforms