Completed DynamoDB + DAX Benchmarker with a nice TUI to boot

2023-08-02 18:11:41 -06:00
parent 09862c1b43
commit e42070eefa
55 changed files with 3574 additions and 1 deletions
@@ -0,0 +1,72 @@
+use aws_sdk_dynamodb::types::AttributeValue;
+use log::{error, info};
+
+use crate::{models::DynamoDbSimulationMetrics, time};
+
+use super::{utils, Simulator};
+
+impl<'a> Simulator<'a> {
+  pub(super) async fn assert_item_was_created(
+    &mut self,
+    id: AttributeValue,
+    metrics: &mut DynamoDbSimulationMetrics,
+  ) -> anyhow::Result<()> {
+    let partition_key = utils::extract_partition_key(id.clone());
+    let mut attempts_exhausted = false;
+
+    let write_confirmation_time = time!(for i in 0..10 {
+      info!("Attempt {i}: Fetching newly added item with partition key: {partition_key}");
+
+      match self.read_item(id.clone(), metrics, false).await? {
+        Some(_) => {
+          info!("Successfully read new item with partition key: {partition_key}");
+          break;
+        }
+        None => {
+          error!("Unable to find new item with partition key: {partition_key}");
+          if i == 9 {
+            error!("All attempts to fetch the newly added item with partition key: {partition_key} failed!");
+            attempts_exhausted = true;
+          }
+        }
+      };
+    });
+
+    if !attempts_exhausted {
+      metrics.write_item_confirmation_time = Some(write_confirmation_time);
+    }
+
+    Ok(())
+  }
+
+  pub(super) async fn assert_item_was_deleted(
+    &mut self,
+    id: AttributeValue,
+    metrics: &mut DynamoDbSimulationMetrics,
+  ) -> anyhow::Result<()> {
+    let partition_key = utils::extract_partition_key(id.clone());
+    let mut attempts_exhausted = false;
+    let delete_confirmation_time = time!(for i in 0..10 {
+      info!("Attempt {i}: Fetching deleted item with partition key: {partition_key}...");
+      match self.read_item(id.clone(), metrics, false).await? {
+        Some(_) => {
+          error!("Item with partition key {partition_key} was not deleted as expected!");
+          if i == 9 {
+            error!("All attempts to receive an empty response to verify item with partition key: {partition_key} was deleted failed!");
+            attempts_exhausted = true;
+          }
+        }
+        None => {
+          info!("Item with partition key {partition_key} was successfully deleted.");
+          break;
+        }
+      }
+    });
+
+    if !attempts_exhausted {
+      metrics.delete_item_confirmation_time = Some(delete_confirmation_time);
+    }
+
+    Ok(())
+  }
+}
@@ -0,0 +1,140 @@
+use aws_sdk_dynamodb::{types::AttributeValue, Client};
+use log::{error, info};
+use rand::{
+  rngs::{OsRng, StdRng},
+  Rng, SeedableRng,
+};
+
+use crate::{models::DynamoDbSimulationMetrics, time};
+
+mod assertions;
+mod operations;
+mod utils;
+
+pub struct Simulator<'a> {
+  dynamodb_client: &'a Client,
+  table_name: String,
+  attributes: u32,
+  partition_keys_vec: &'a [String],
+  rng: StdRng,
+}
+
+impl<'a> Simulator<'a> {
+  pub fn new(
+    dynamodb_client: &'a Client,
+    table_name: String,
+    attributes: u32,
+    partition_keys_vec: &'a [String],
+  ) -> Simulator<'a> {
+    Simulator {
+      dynamodb_client,
+      table_name,
+      attributes,
+      partition_keys_vec,
+      rng: StdRng::from_seed(OsRng.gen()),
+    }
+  }
+
+  pub async fn simulate_read_operation(
+    &mut self,
+    metrics: &mut DynamoDbSimulationMetrics,
+  ) -> anyhow::Result<()> {
+    info!("Performing READ Operation...");
+    let partition_key =
+      self.partition_keys_vec[self.rng.gen_range(0..self.partition_keys_vec.len())].clone();
+    let id = AttributeValue::S(partition_key.clone());
+
+    for i in 0..10 {
+      info!("Attempt {i}: Fetching existing item with partition key: {partition_key}");
+
+      match self.read_item(id.clone(), metrics, true).await? {
+        Some(_) => {
+          info!("Successfully read existing item with partition key: {partition_key}");
+          break;
+        }
+        None => {
+          error!("Unable to find existing item with partition key: {partition_key}");
+          if i == 9 {
+            error!(
+              "All attempts to fetch the existing item with partition key: {partition_key} failed!"
+            );
+          }
+        }
+      }
+    }
+
+    Ok(())
+  }
+
+  pub async fn simulate_write_operation(
+    &mut self,
+    metrics: &mut DynamoDbSimulationMetrics,
+  ) -> anyhow::Result<()> {
+    info!("Performing WRITE operation...");
+    let benchmarking_item = self.put_item(metrics).await?;
+    let id = benchmarking_item.get_id();
+
+    self.assert_item_was_created(id.clone(), metrics).await?;
+
+    self.delete_item(id.clone(), metrics).await?;
+
+    self.assert_item_was_deleted(id, metrics).await?;
+
+    Ok(())
+  }
+
+  pub async fn simulate_update_operation(
+    &mut self,
+    metrics: &mut DynamoDbSimulationMetrics,
+  ) -> anyhow::Result<()> {
+    info!("Performing UPDATE operation...");
+    let new_item = self.put_item(metrics).await?;
+    let id = new_item.get_id();
+    let partition_key = utils::extract_partition_key(id.clone());
+    let mut attempts_exhausted = false;
+
+    self.assert_item_was_created(id.clone(), metrics).await?;
+    self.update_item(id.clone(), metrics).await?;
+
+    let update_confirmation_time = time!(for i in 0..10 {
+      info!("Attempt {i}: Fetching updated item for partition key: {partition_key}...");
+
+      let updated_item = self.read_item(id.clone(), metrics, false).await?.unwrap();
+
+      let new_item_attribute_value = new_item
+        .get("1")
+        .cloned()
+        .unwrap()
+        .as_n()
+        .unwrap()
+        .to_string();
+      let updated_item_attribute_value = updated_item
+        .get("1")
+        .cloned()
+        .unwrap()
+        .as_n()
+        .unwrap()
+        .to_string();
+
+      if new_item_attribute_value != updated_item_attribute_value {
+        info!("Confirmed update for partition key: {partition_key}");
+        break;
+      } else {
+        error!("Update for partition key {partition_key} failed! Values are still equal!");
+        if i == 9 {
+          error!("Exhausted attempts to fetch updated item!");
+          attempts_exhausted = true;
+        }
+      }
+    });
+
+    if !attempts_exhausted {
+      metrics.update_item_confirmation_time = Some(update_confirmation_time);
+    }
+
+    self.delete_item(id.clone(), metrics).await?;
+    self.assert_item_was_deleted(id, metrics).await?;
+
+    Ok(())
+  }
+}
@@ -0,0 +1,144 @@
+use anyhow::anyhow;
+use aws_sdk_dynamodb::types::AttributeValue;
+use log::{error, info};
+
+use crate::{
+  models::{BenchmarkingItem, DynamoDbSimulationMetrics},
+  time,
+};
+
+use super::{utils::extract_partition_key, Simulator};
+
+impl<'a> Simulator<'a> {
+  pub async fn read_item(
+    &mut self,
+    id: AttributeValue,
+    metrics: &mut DynamoDbSimulationMetrics,
+    record_metrics: bool,
+  ) -> anyhow::Result<Option<BenchmarkingItem>> {
+    let partition_key = extract_partition_key(id.clone());
+    let (read_time, response) = time!(
+      resp,
+      self
+        .dynamodb_client
+        .get_item()
+        .table_name(self.table_name.clone())
+        .key("id", id)
+        .send()
+        .await
+    );
+
+    if record_metrics {
+      metrics.read_time = Some(read_time);
+    }
+
+    match response {
+      Ok(resp) => {
+        info!("Found item: {}", partition_key);
+        if let Some(item) = resp.item() {
+          info!("Fetched item: {item:?}");
+          Ok(Some(BenchmarkingItem::from(item.clone())))
+        } else {
+          info!("No items found with partition key: {partition_key}");
+          Ok(None)
+        }
+      }
+      Err(e) => {
+        error!("Could not fetch item with partition key: {partition_key}. {e:?}");
+        Err(anyhow!(e))
+      }
+    }
+  }
+
+  pub async fn update_item(
+    &mut self,
+    id: AttributeValue,
+    metrics: &mut DynamoDbSimulationMetrics,
+  ) -> anyhow::Result<()> {
+    let mut updated_item = BenchmarkingItem::new(self.attributes);
+    updated_item.insert("id", id.clone());
+    let partition_key = extract_partition_key(id);
+    let (update_time, response) = time!(
+      resp,
+      self
+        .dynamodb_client
+        .put_item()
+        .table_name(self.table_name.clone())
+        .set_item(Some(updated_item.extract_map()))
+        .send()
+        .await
+    );
+    metrics.update_time = Some(update_time);
+
+    match response {
+      Ok(_) => {
+        info!("Successfully updated item with partition_key: {partition_key}");
+        Ok(())
+      }
+      Err(e) => {
+        error!("Could not update item with partition key: {partition_key}. {e:?}");
+        Err(anyhow!(e))
+      }
+    }
+  }
+
+  pub async fn put_item(
+    &mut self,
+    metrics: &mut DynamoDbSimulationMetrics,
+  ) -> anyhow::Result<BenchmarkingItem> {
+    let new_item = BenchmarkingItem::new(self.attributes);
+    let partition_key = extract_partition_key(new_item.get("id").cloned().unwrap());
+    let (time, response) = time!(
+      resp,
+      self
+        .dynamodb_client
+        .put_item()
+        .table_name(self.table_name.clone())
+        .set_item(Some(new_item.extract_map()))
+        .send()
+        .await
+    );
+    metrics.write_time = Some(time);
+
+    match response {
+      Ok(_) => {
+        info!("Successfully put new item with partition key: {partition_key}");
+        Ok(new_item)
+      }
+      Err(e) => {
+        error!("Could not put new item with partition key: {partition_key}. {e:?}");
+        Err(anyhow!(e))
+      }
+    }
+  }
+
+  pub async fn delete_item(
+    &mut self,
+    id: AttributeValue,
+    metrics: &mut DynamoDbSimulationMetrics,
+  ) -> anyhow::Result<()> {
+    let partition_key = extract_partition_key(id.clone());
+    let (delete_time, response) = time!(
+      resp,
+      self
+        .dynamodb_client
+        .delete_item()
+        .table_name(self.table_name.clone())
+        .key("id", id)
+        .send()
+        .await
+    );
+    metrics.delete_time = Some(delete_time);
+
+    match response {
+      Ok(_) => {
+        info!("Successfully deleted item with partition key: {partition_key}");
+        Ok(())
+      }
+      Err(e) => {
+        error!("Could not delete item with partition key: {partition_key}. {e:?}");
+        Err(anyhow!(e))
+      }
+    }
+  }
+}
@@ -0,0 +1,5 @@
+use aws_sdk_dynamodb::types::AttributeValue;
+
+pub(super) fn extract_partition_key(id: AttributeValue) -> String {
+  id.clone().as_s().unwrap().to_string()
+}