Update serializeError to allow with and without envelope

Summary: Provide serializeError methods that can be templated on whether or not we wish to include the envelope. Reviewed By: sethdelliott Differential Revision: D28976217 fbshipit-source-id: 4fd14bbd4bccab691b2ac4b99fe70ab2e1b82924
facebook · Jul 9, 2021 · b0a7a48 · b0a7a48
1 parent d0d6cde
commit b0a7a48
Show file tree

Hide file tree

Showing 2 changed files with 79 additions and 61 deletions.
diff --git a/thrift/lib/cpp2/protocol/Serializer.cpp b/thrift/lib/cpp2/protocol/Serializer.cpp
@@ -19,48 +19,6 @@
 namespace apache {
 namespace thrift {
 
-std::unique_ptr<folly::IOBuf> serializeError(
-    int protId, const TApplicationException& obj, folly::IOBuf* buf) {
-  switch (protId) {
-    case apache::thrift::protocol::T_BINARY_PROTOCOL: {
-      return serializeErrorProtocol<BinaryProtocolReader, BinaryProtocolWriter>(
-          obj, std::move(buf));
-    }
-    case apache::thrift::protocol::T_COMPACT_PROTOCOL: {
-      return serializeErrorProtocol<
-          CompactProtocolReader,
-          CompactProtocolWriter>(obj, std::move(buf));
-    }
-    default: {
-      LOG(ERROR) << "Invalid protocol from client";
-    }
-  }
-
-  return nullptr;
-}
-
-std::unique_ptr<folly::IOBuf> serializeError(
-    int protId,
-    const TApplicationException& obj,
-    const std::string& fname,
-    int32_t protoSeqId) {
-  switch (protId) {
-    case apache::thrift::protocol::T_BINARY_PROTOCOL: {
-      return serializeErrorProtocol<BinaryProtocolWriter>(
-          obj, fname, protoSeqId);
-    }
-    case apache::thrift::protocol::T_COMPACT_PROTOCOL: {
-      return serializeErrorProtocol<CompactProtocolWriter>(
-          obj, fname, protoSeqId);
-    }
-    default: {
-      LOG(ERROR) << "Invalid protocol from client";
-    }
-  }
-
-  return nullptr;
-}
-
 std::unique_ptr<folly::IOBuf> serializeErrorStruct(
     protocol::PROTOCOL_TYPES protId, const TApplicationException& obj) {
   auto f = [&](auto prot) -> std::unique_ptr<folly::IOBuf> {

diff --git a/thrift/lib/cpp2/protocol/Serializer.h b/thrift/lib/cpp2/protocol/Serializer.h
@@ -216,51 +216,111 @@ typedef Serializer<SimpleJSONProtocolReader, SimpleJSONProtocolWriter>
 typedef Serializer<NimbleProtocolReader, NimbleProtocolWriter> NimbleSerializer;
 
 // Serialization code specific to handling errors
-template <typename ProtIn, typename ProtOut>
+template <typename ProtIn, typename ProtOut, bool includeEnvelope = true>
 std::unique_ptr<folly::IOBuf> serializeErrorProtocol(
     const TApplicationException& obj, folly::IOBuf* req) {
-  ProtIn iprot;
-  std::string fname;
-  apache::thrift::MessageType mtype;
-  int32_t protoSeqId = 0;
-  iprot.setInput(req);
-  iprot.readMessageBegin(fname, mtype, protoSeqId);
-
   ProtOut prot;
-  size_t bufSize = obj.serializedSizeZC(&prot);
-  bufSize += prot.serializedMessageSize(fname);
   folly::IOBufQueue queue;
-  prot.setOutput(&queue, bufSize);
-  prot.writeMessageBegin(fname, MessageType::T_EXCEPTION, protoSeqId);
+  size_t objSize = obj.serializedSizeZC(&prot);
+  if /*constexpr*/ (includeEnvelope) {
+    ProtIn iprot;
+    std::string fname;
+    apache::thrift::MessageType mtype;
+    int32_t protoSeqId = 0;
+    iprot.setInput(req);
+    iprot.readMessageBegin(fname, mtype, protoSeqId);
+    prot.setOutput(&queue, objSize + prot.serializedMessageSize(fname));
+    prot.writeMessageBegin(fname, MessageType::T_EXCEPTION, protoSeqId);
+  } else {
+    prot.setOutput(&queue, objSize);
+  }
   obj.write(&prot);
   prot.writeMessageEnd();
   return queue.move();
 }
 
-template <typename ProtOut>
+template <typename ProtOut, bool includeEnvelope = true>
 std::unique_ptr<folly::IOBuf> serializeErrorProtocol(
     const TApplicationException& obj,
     const std::string& fname,
     int32_t protoSeqId) {
   ProtOut prot;
-  size_t bufSize = obj.serializedSizeZC(&prot);
-  bufSize += prot.serializedMessageSize(fname);
   folly::IOBufQueue queue;
-  prot.setOutput(&queue, bufSize);
-  prot.writeMessageBegin(fname, MessageType::T_EXCEPTION, protoSeqId);
+  std::size_t objSize = obj.serializedSizeZC(&prot);
+  if /*constexpr*/ (includeEnvelope) {
+    prot.setOutput(&queue, objSize + prot.serializedMessageSize(fname));
+    prot.writeMessageBegin(fname, MessageType::T_EXCEPTION, protoSeqId);
+  } else {
+    prot.setOutput(&queue, objSize);
+  }
   obj.write(&prot);
   prot.writeMessageEnd();
   return queue.move();
 }
 
+template <bool includeEnvelope = true>
 std::unique_ptr<folly::IOBuf> serializeError(
-    int protId, const TApplicationException& obj, folly::IOBuf* buf);
+    int protId, const TApplicationException& obj, folly::IOBuf* buf) {
+  switch (protId) {
+    case apache::thrift::protocol::T_BINARY_PROTOCOL: {
+      return serializeErrorProtocol<
+          BinaryProtocolReader,
+          BinaryProtocolWriter,
+          includeEnvelope>(obj, buf);
+    }
+    case apache::thrift::protocol::T_COMPACT_PROTOCOL: {
+      return serializeErrorProtocol<
+          CompactProtocolReader,
+          CompactProtocolWriter,
+          includeEnvelope>(obj, buf);
+    }
+    default: {
+      LOG(ERROR) << "Invalid protocol from client";
+    }
+  }
 
+  return nullptr;
+}
+
+template <bool includeEnvelope = true>
 std::unique_ptr<folly::IOBuf> serializeError(
     int protId,
     const TApplicationException& obj,
     const std::string& fname,
-    int32_t protoSeqId);
+    int32_t protoSeqId) {
+  switch (protId) {
+    case apache::thrift::protocol::T_BINARY_PROTOCOL: {
+      return serializeErrorProtocol<BinaryProtocolWriter, includeEnvelope>(
+          obj, fname, protoSeqId);
+    }
+    case apache::thrift::protocol::T_COMPACT_PROTOCOL: {
+      return serializeErrorProtocol<CompactProtocolWriter, includeEnvelope>(
+          obj, fname, protoSeqId);
+    }
+    default: {
+      LOG(ERROR) << "Invalid protocol from client";
+    }
+  }
+
+  return nullptr;
+}
+
+// For places where we can't currently invoke the templated version directly,
+inline std::unique_ptr<folly::IOBuf> serializeErrorWithEnvelope(
+    int protId,
+    const TApplicationException& obj,
+    const std::string& fname,
+    int32_t protoSeqId) {
+  return serializeError<true>(protId, obj, fname, protoSeqId);
+}
+
+inline std::unique_ptr<folly::IOBuf> serializeErrorWithoutEnvelope(
+    int protId,
+    const TApplicationException& obj,
+    const std::string& fname,
+    int32_t protoSeqId) {
+  return serializeError<false>(protId, obj, fname, protoSeqId);
+}
 
 // serialize TApplicationException without a protocol message envelop
 std::unique_ptr<folly::IOBuf> serializeErrorStruct(