intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

1 files changed, 90 insertions, 0 deletions

diff --git a/library/cpp/grpc/server/grpc_response.h b/library/cpp/grpc/server/grpc_response.h
new file mode 100644
index 0000000000..8e9afe44d5
--- /dev/null
+++ b/library/cpp/grpc/server/grpc_response.h
@@ -0,0 +1,90 @@
+#pragma once
+
+#include <grpc++/impl/codegen/byte_buffer.h>
+#include <grpc++/impl/codegen/proto_utils.h>
+
+#include <variant>
+
+namespace NGrpc {
+
+/**
+ * Universal response that owns underlying message or buffer.
+ */
+template <typename TMsg>
+class TUniversalResponse: public TAtomicRefCount<TUniversalResponse<TMsg>>, public TMoveOnly {
+    friend class grpc::SerializationTraits<TUniversalResponse<TMsg>>;
+
+public:
+    explicit TUniversalResponse(NProtoBuf::Message* msg) noexcept
+        : Data_{TMsg{}}
+    {
+        std::get<TMsg>(Data_).Swap(static_cast<TMsg*>(msg));
+    }
+
+    explicit TUniversalResponse(grpc::ByteBuffer* buffer) noexcept
+        : Data_{grpc::ByteBuffer{}}
+    {
+        std::get<grpc::ByteBuffer>(Data_).Swap(buffer);
+    }
+
+private:
+    std::variant<TMsg, grpc::ByteBuffer> Data_;
+};
+
+/**
+ * Universal response that only keeps reference to underlying message or buffer.
+ */
+template <typename TMsg>
+class TUniversalResponseRef: private TMoveOnly {
+    friend class grpc::SerializationTraits<TUniversalResponseRef<TMsg>>;
+
+public:
+    explicit TUniversalResponseRef(const NProtoBuf::Message* msg)
+        : Data_{msg}
+    {
+    }
+
+    explicit TUniversalResponseRef(const grpc::ByteBuffer* buffer)
+        : Data_{buffer}
+    {
+    }
+
+private:
+    std::variant<const NProtoBuf::Message*, const grpc::ByteBuffer*> Data_;
+};
+
+} // namespace NGrpc
+
+namespace grpc {
+
+template <typename TMsg>
+class SerializationTraits<NGrpc::TUniversalResponse<TMsg>> {
+public:
+    static Status Serialize(
+            const NGrpc::TUniversalResponse<TMsg>& resp,
+            ByteBuffer* buffer,
+            bool* ownBuffer)
+    {
+        return std::visit([&](const auto& data) {
+            using T = std::decay_t<decltype(data)>;
+            return SerializationTraits<T>::Serialize(data, buffer, ownBuffer);
+        }, resp.Data_);
+    }
+};
+
+template <typename TMsg>
+class SerializationTraits<NGrpc::TUniversalResponseRef<TMsg>> {
+public:
+    static Status Serialize(
+        const NGrpc::TUniversalResponseRef<TMsg>& resp,
+        ByteBuffer* buffer,
+        bool* ownBuffer)
+    {
+        return std::visit([&](const auto* data) {
+            using T = std::decay_t<std::remove_pointer_t<decltype(data)>>;
+            return SerializationTraits<T>::Serialize(*data, buffer, ownBuffer);
+        }, resp.Data_);
+    }
+};
+
+} // namespace grpc