Kafka codec: precompute request size before serialization, so we do n…

source/extensions/filters/network/kafka/kafka_request.h

@@ -57,7 +57,13 @@ class AbstractRequest {
   AbstractRequest(const RequestHeader& request_header) : request_header_{request_header} {};
 
   /**
-   * Encode the contents of this message into a given buffer.
+   * Computes the size of this request, if it were to be serialized.
+   * @return serialized size of request
+   */
+  virtual size_t computeSize() const PURE;
+
+  /**
+   * Encode the contents of this request into a given buffer.
    * @param dst buffer instance to keep serialized message
    */
   virtual size_t encode(Buffer::Instance& dst) const PURE;
@@ -82,6 +88,22 @@ template <typename Data> class Request : public AbstractRequest {
   Request(const RequestHeader& request_header, const Data& data)
       : AbstractRequest{request_header}, data_{data} {};
 
+  /**
+   * Compute the size of request, what includes both the request header and its real data.
+   */
+  size_t computeSize() const override {
+    const EncodingContext context{request_header_.api_version_};
+    size_t result = 0;
+    // Compute size of header.
+    result += context.computeSize(request_header_.api_key_);
+    result += context.computeSize(request_header_.api_version_);
+    result += context.computeSize(request_header_.correlation_id_);
+    result += context.computeSize(request_header_.client_id_);
+    // Compute size of request data.
+    result += context.computeSize(data_);
+    return result;
+  }
+
   /**
    * Encodes given request into a buffer, with any extra configuration carried by the context.
    */

source/extensions/filters/network/kafka/protocol_code_generator/complex_type_template.j2

@@ -22,6 +22,23 @@ struct {{ complex_type.name }} {
   // constructor used in versions: {{ constructor['versions'] }}
   {{ constructor['full_declaration'] }}{% endfor %}
 
+  {# For every field that's used in version, just compute its size using an encoder. #}
+  {% if complex_type.fields|length > 0 %}
+  size_t computeSize(const EncodingContext& encoder) const {
+    const int16_t api_version = encoder.apiVersion();
+    size_t written{0};{% for field in complex_type.fields %}
+    if (api_version >= {{ field.version_usage[0] }}
+      && api_version < {{ field.version_usage[-1] + 1 }}) {
+      written += encoder.computeSize({{ field.name }}_);
+    }{% endfor %}
+    return written;
+  }
+  {% else %}
+  size_t computeSize(const EncodingContext&) const {
+    return 0;
+  }
+  {% endif %}
+
   {# For every field that's used in version, just serialize it. #}
   {% if complex_type.fields|length > 0 %}
   size_t encode(Buffer::Instance& dst, EncodingContext& encoder) const {

source/extensions/filters/network/kafka/request_codec.cc

@@ -80,14 +80,9 @@ void RequestDecoder::doParse(const Buffer::RawSlice& slice) {
 }
 
 void RequestEncoder::encode(const AbstractRequest& message) {
-  Buffer::OwnedImpl data_buffer;
-  // TODO(adamkotwasinski) Precompute the size instead of using temporary buffer.
-  // When we have the 'computeSize' method, then we can push encoding request's size into
-  // Request::encode
-  int32_t data_len = message.encode(data_buffer); // Encode data and compute data length.
-  EncodingContext encoder{-1};
-  encoder.encode(data_len, output_); // Encode data length into result.
-  output_.add(data_buffer);          // Copy encoded data into result.
+  const int32_t size = htobe32(message.computeSize());
+  output_.add(&size, sizeof(size)); // Encode data length.
+  message.encode(output_);          // Encode data.
 }
 
 } // namespace Kafka

source/extensions/filters/network/kafka/serialization.h

@@ -469,6 +469,24 @@ class EncodingContext {
 public:
   EncodingContext(int16_t api_version) : api_version_{api_version} {};
 
+  /**
+   * Compute size of given reference, if it were to be encoded.
+   * @return serialized size of argument.
+   */
+  template <typename T> size_t computeSize(const T& arg) const;
+
+  /**
+   * Compute size of given array, if it were to be encoded.
+   * @return serialized size of argument.
+   */
+  template <typename T> size_t computeSize(const std::vector<T>& arg) const;
+
+  /**
+   * Compute size of given nullable array, if it were to be encoded.
+   * @return serialized size of argument.
+   */
+  template <typename T> size_t computeSize(const NullableArray<T>& arg) const;
+
   /**
    * Encode given reference in a buffer.
    * @return bytes written
@@ -493,6 +511,82 @@ class EncodingContext {
   const int16_t api_version_;
 };
 
+/**
+ * For non-primitive types, call `computeSize` on them, to delegate the work to the entity itself.
+ * The entity may use the information in context to decide which fields are included etc.
+ */
+template <typename T> inline size_t EncodingContext::computeSize(const T& arg) const {
+  return arg.computeSize(*this);
+}
+
+/**
+ * For primitive types, Kafka size == sizeof(x).
+ */
+#define COMPUTE_SIZE_OF_NUMERIC_TYPE(TYPE)                                                         \
+  template <> constexpr size_t EncodingContext::computeSize(const TYPE&) const {                   \
+    return sizeof(TYPE);                                                                           \
+  }
+
+COMPUTE_SIZE_OF_NUMERIC_TYPE(bool)
+COMPUTE_SIZE_OF_NUMERIC_TYPE(int8_t)
+COMPUTE_SIZE_OF_NUMERIC_TYPE(int16_t)
+COMPUTE_SIZE_OF_NUMERIC_TYPE(int32_t)
+COMPUTE_SIZE_OF_NUMERIC_TYPE(uint32_t)
+COMPUTE_SIZE_OF_NUMERIC_TYPE(int64_t)
+
+/**
+ * Template overload for string.
+ * Kafka String's size is INT16 for header + N bytes.
+ */
+template <> inline size_t EncodingContext::computeSize(const std::string& arg) const {
+  return sizeof(int16_t) + arg.size();
+}
+
+/**
+ * Template overload for nullable string.
+ * Kafka NullableString's size is INT16 for header + N bytes (N >= 0).
+ */
+template <> inline size_t EncodingContext::computeSize(const NullableString& arg) const {
+  return sizeof(int16_t) + (arg ? arg->size() : 0);
+}
+
+/**
+ * Template overload for byte array.
+ * Kafka byte array size is INT32 for header + N bytes.
+ */
+template <> inline size_t EncodingContext::computeSize(const Bytes& arg) const {
+  return sizeof(int32_t) + arg.size();
+}
+
+/**
+ * Template overload for nullable byte array.
+ * Kafka nullable byte array size is INT32 for header + N bytes (N >= 0).
+ */
+template <> inline size_t EncodingContext::computeSize(const NullableBytes& arg) const {
+  return sizeof(int32_t) + (arg ? arg->size() : 0);
+}
+
+/**
+ * Template overload for Array of T.
+ * The size of array is size of header and all of its elements.
+ */
+template <typename T> inline size_t EncodingContext::computeSize(const std::vector<T>& arg) const {
+  size_t result = sizeof(int32_t);
+  for (const T& el : arg) {
+    result += computeSize(el);
+  }
+  return result;
+}
+
+/**
+ * Template overload for NullableArray of T.
+ * The size of array is size of header and all of its elements.
+ */
+template <typename T>
+inline size_t EncodingContext::computeSize(const NullableArray<T>& arg) const {
+  return arg ? computeSize(*arg) : sizeof(int32_t);
+}
+
 /**
  * For non-primitive types, call `encode` on them, to delegate the serialization to the entity
  * itself.