DynamoDB: Log stats for partition id per operation

source/common/dynamo/dynamo_filter.cc

@@ -57,6 +57,8 @@ void DynamoFilter::onEncodeComplete(const Buffer::Instance& data) {
 
     chargeBasicStats(status);
 
+    chargeTablePartitionIdStats(data);
+
     if (Http::CodeUtility::is4xx(status)) {
       chargeFailureSpecificStats(data);
     }
@@ -208,4 +210,44 @@ void DynamoFilter::chargeFailureSpecificStats(const Buffer::Instance& data) {
   }
 }
 
+void DynamoFilter::chargeTablePartitionIdStats(const Buffer::Instance& data) {
+  // Only log partition stats for single table operations.
+  if (table_descriptor_.table_name.empty() || operation_.empty()) {
+    return;
+  }
+  std::string body = buildBody(encoder_callbacks_->encodingBuffer(), data);
+  if (!body.empty()) {
+    try {
+      std::vector<Dynamo::RequestParser::PartitionDescriptor> partitions =
+          Dynamo::RequestParser::parsePartitions(body);
+      for (const Dynamo::RequestParser::PartitionDescriptor& partition : partitions) {
+        std::string stats_string = DynamoFilter::buildPartitionStatString(
+            stat_prefix_, table_descriptor_.table_name, operation_, partition.partition_id_);
+        stats_.counter(stats_string).add(partition.capacity_);
+      }
+    } catch (const Json::Exception& jsonEx) {
+      // Body parsing failed. This should not happen, just put a stat for that.
+      stats_.counter(fmt::format("{}invalid_resp_body", stat_prefix_)).inc();
+    }
+  }
+}
+
+std::string DynamoFilter::buildPartitionStatString(const std::string& stat_prefix,
+                                                   const std::string& table_name,
+                                                   const std::string& operation,
+                                                   const std::string& partition_id) {
+  std::string stats_table_prefix = fmt::format("{}table.{}", stat_prefix, table_name);
+  // Use only the last 7 characters from the partition id.
+  std::string stats_partition_postfix =
+      fmt::format(".Capacity.{}.__partition_id={}", operation,
+                  partition_id.substr(partition_id.size() - 7, partition_id.size()));
+  // Calculate how many characters are available for the table prefix.
+  size_t remaining_size = MAX_NAME_SIZE - stats_partition_postfix.size() - 1;
+  // Truncate the table prefix if the curent string is too large.
+  if (stats_table_prefix.size() > remaining_size) {
+    stats_table_prefix = stats_table_prefix.substr(0, remaining_size);
+  }
+  return fmt::format("{}{}", stats_table_prefix, stats_partition_postfix);
+}
+
 } // Dynamo

source/common/dynamo/dynamo_filter.h

@@ -37,6 +37,11 @@ class DynamoFilter : public Http::StreamFilter {
     encoder_callbacks_ = &callbacks;
   }
 
+  static std::string buildPartitionStatString(const std::string& stat_prefix,
+                                              const std::string& table_name,
+                                              const std::string& operation,
+                                              const std::string& partition_id);
+
 private:
   void onDecodeComplete(const Buffer::Instance& data);
   void onEncodeComplete(const Buffer::Instance& data);
@@ -46,11 +51,13 @@ class DynamoFilter : public Http::StreamFilter {
                             uint64_t status);
   void chargeFailureSpecificStats(const Buffer::Instance& data);
   void chargeUnProcessedKeysStats(const Buffer::Instance& data);
+  void chargeTablePartitionIdStats(const Buffer::Instance& data);
+
+  static const size_t MAX_NAME_SIZE = 127;
 
   Runtime::Loader& runtime_;
   std::string stat_prefix_;
   Stats::Store& stats_;
-
   bool enabled_{};
   std::string operation_{};
   RequestParser::TableDescriptor table_descriptor_{"", true};

source/common/dynamo/dynamo_request_parser.cc

@@ -125,4 +125,28 @@ bool RequestParser::isBatchOperation(const std::string& operation) {
          BATCH_OPERATIONS.end();
 }
 
+std::vector<RequestParser::PartitionDescriptor>
+RequestParser::parsePartitions(const std::string& data) {
+  Json::StringLoader json(data);
+  std::vector<RequestParser::PartitionDescriptor> partition_descriptors;
+
+  if (json.hasObject("ConsumedCapacity")) {
+    Json::Object consumed_capacity = json.getObject("ConsumedCapacity");
+    if (consumed_capacity.hasObject("Partitions")) {
+      Json::Object partitions = consumed_capacity.getObject("Partitions");
+
+      partitions.iterate(
+          [&partition_descriptors, &partitions](const std::string& key, const Json::Object&) {
+            // Capacity is a double and it is rounded up to the nearest integer.
+            // The partition stat will be incremented by the capacity value.
+            uint64_t capacity_integer =
+                static_cast<uint64_t>(std::ceil(partitions.getDouble(key, 0.0)));
+            partition_descriptors.emplace_back(key, capacity_integer);
+            return true;
+          });
+    }
+  }
+  return partition_descriptors;
+}
+
 } // Dynamo

source/common/dynamo/dynamo_request_parser.h

@@ -18,6 +18,13 @@ class RequestParser {
     bool is_single_table;
   };
 
+  struct PartitionDescriptor {
+    PartitionDescriptor(std::string partition, uint64_t capacity)
+        : partition_id_(partition), capacity_(capacity) {}
+    std::string partition_id_;
+    uint64_t capacity_;
+  };
+
   /**
    * Parse operation out of x-amz-target header.
    * @return empty string if operation cannot be parsed.
@@ -66,6 +73,12 @@ class RequestParser {
    */
   static bool isBatchOperation(const std::string& operation);
 
+  /**
+   * Parse the Partition ids and the capacity from the response body.
+   * @return empty set if there are no partition ids or a set of partition ids
+   */
+  static std::vector<PartitionDescriptor> parsePartitions(const std::string& data);
+
 private:
   static const Http::LowerCaseString X_AMZ_TARGET;
   static const std::vector<std::string> SINGLE_TABLE_OPERATIONS;

source/common/json/json_loader.cc

@@ -136,6 +136,23 @@ std::vector<std::string> Object::getStringArray(const std::string& name) const {
   return string_array;
 }
 
+double Object::getDouble(const std::string& name) const {
+  json_t* real = json_object_get(json_, name.c_str());
+  if (!real || !json_is_real(real)) {
+    throw Exception(fmt::format("key '{}' missing or not an double in '{}'", name, name_));
+  }
+
+  return json_real_value(real);
+}
+
+double Object::getDouble(const std::string& name, const double& default_value) const {
+  if (!json_object_get(json_, name.c_str())) {
+    return default_value;
+  } else {
+    return getDouble(name);
+  }
+}
+
 void Object::iterate(const ObjectCallback& callback) {
   const char* key;
   json_t* value;

source/common/json/json_loader.h

@@ -81,14 +81,14 @@ class Object {
 
   /**
    * Get a string value by name.
-   * @param name suppies the key name.
+   * @param name supplies the key name.
    * @return std::string the value.
    */
   std::string getString(const std::string& name) const;
 
   /**
    * Get a string value by name or return a default if name does not exist.
-   * @param name suppies the key name.
+   * @param name supplies the key name.
    * @param default_value supplies the value to return if name does not exist.
    * @return std::string the value.
    */
@@ -101,6 +101,21 @@ class Object {
    */
   std::vector<std::string> getStringArray(const std::string& name) const;
 
+  /**
+   * Get a double value by name.
+   * @param name supplies the key name.
+   * @return double the value.
+   */
+  double getDouble(const std::string& name) const;
+
+  /**
+   * Get a double value by name.
+   * @param name supplies the key name.
+   * @param default_value supplies the value to return if name does not exist.
+   * @return double the value.
+   */
+  double getDouble(const std::string& name, const double& default_value) const;
+
   /**
    * Iterate Object and call callback on key-value pairs
    */

source/precompiled/precompiled.h

@@ -2,6 +2,7 @@
 
 #include <arpa/inet.h>
 #include <chrono>
+#include <cmath>
 #include <condition_variable>
 #include <fstream>
 #include <iostream>