From 0481ed565ff3286e7fdaa488d2874b519772b405 Mon Sep 17 00:00:00 2001
From: Jake <j82w@users.noreply.github.com>
Date: Sun, 6 Oct 2024 01:18:52 -0400
Subject: [PATCH] Benchmark: Adds a workload to insert and read the inserted
 item. (#4482)

# Pull Request Template

## Description

Adds a new benchmark to simulate users inserting a document and
immediately reading it. This with the
https://github.com/Azure/azure-cosmos-dotnet-v3/pull/4478 will test
replication latency between different geo regions, and how SDK handles
various error scenarios.

## Type of change

Please delete options that are not relevant.

- [] Bug fix (non-breaking change which fixes an issue)
- [x] New feature (non-breaking change which adds functionality)
- [] Breaking change (fix or feature that would cause existing
functionality to not work as expected)
- [] This change requires a documentation update

## Closing issues

To automatically close an issue: closes #IssueNumber

Co-authored-by: jakewilley_microsoft <--global>
Co-authored-by: Kiran Kumar Kolli <kirankk@microsoft.com>
---
 .../v3/InsertReadV3BenchmarkOperation.cs      | 97 +++++++++++++++++++
 1 file changed, 97 insertions(+)
 create mode 100644 Microsoft.Azure.Cosmos.Samples/Tools/Benchmark/v3/InsertReadV3BenchmarkOperation.cs

diff --git a/Microsoft.Azure.Cosmos.Samples/Tools/Benchmark/v3/InsertReadV3BenchmarkOperation.cs b/Microsoft.Azure.Cosmos.Samples/Tools/Benchmark/v3/InsertReadV3BenchmarkOperation.cs
new file mode 100644
index 0000000000..d4a5cd4ce8
--- /dev/null
+++ b/Microsoft.Azure.Cosmos.Samples/Tools/Benchmark/v3/InsertReadV3BenchmarkOperation.cs
@@ -0,0 +1,97 @@
+//------------------------------------------------------------
+// Copyright (c) Microsoft Corporation.  All rights reserved.
+//------------------------------------------------------------
+
+namespace CosmosBenchmark
+{
+    using System;
+    using System.Collections.Generic;
+    using System.IO;
+    using System.Linq;
+    using System.Threading.Tasks;
+    using Microsoft.Azure.Cosmos;
+    using Newtonsoft.Json;
+    using Newtonsoft.Json.Serialization;
+
+    internal class InsertReadV3BenchmarkOperation : IBenchmarkOperation
+    {
+        private readonly CosmosClient client;
+        private readonly Container container;
+        private readonly string partitionKeyPath;
+        private readonly Dictionary<string, object> sampleJObject;
+        private readonly string databaseName;
+        private readonly string containerName;
+
+        public InsertReadV3BenchmarkOperation(
+            CosmosClient cosmosClient,
+            string dbName,
+            string containerName,
+            string partitionKeyPath,
+            string sampleJson)
+        {
+            this.databaseName = dbName;
+            this.containerName = containerName;
+            this.client = cosmosClient;
+
+            this.container = cosmosClient.GetContainer(this.databaseName, this.containerName);
+            this.partitionKeyPath = partitionKeyPath.Replace("/", "");
+
+            this.sampleJObject = JsonHelper.Deserialize<Dictionary<string, object>>(sampleJson);
+        }
+
+        public BenchmarkOperationType OperationType => BenchmarkOperationType.Insert;
+
+        public async Task<OperationResult> ExecuteOnceAsync()
+        {
+            PartitionKey partitionKey = new PartitionKey(this.sampleJObject[this.partitionKeyPath].ToString());
+            double ruCharges = 0;
+            CosmosDiagnostics writeDiagnostics = null;
+            using (MemoryStream input = JsonHelper.ToStream(this.sampleJObject))
+            {
+                ResponseMessage itemResponse = await this.container.CreateItemStreamAsync(
+                        input,
+                        partitionKey);
+
+                ruCharges = itemResponse.Headers.RequestCharge;
+                writeDiagnostics = itemResponse.Diagnostics;
+                System.Buffers.ArrayPool<byte>.Shared.Return(input.GetBuffer());
+                
+                if (!itemResponse.IsSuccessStatusCode)
+                {
+                    return new OperationResult()
+                    {
+                        DatabseName = this.databaseName,
+                        ContainerName = this.containerName,
+                        OperationType = this.OperationType,
+                        RuCharges = ruCharges,
+                        CosmosDiagnostics = itemResponse.Diagnostics,
+                        LazyDiagnostics = () => itemResponse.Diagnostics.ToString(),
+                    };
+                }
+            }
+
+            using ResponseMessage readResponse = await this.container.ReadItemStreamAsync(
+                    this.sampleJObject["id"].ToString(),
+                    partitionKey);
+            ruCharges += readResponse.Headers.RequestCharge;
+
+            return new OperationResult()
+            {
+                DatabseName = this.databaseName,
+                ContainerName = this.containerName,
+                OperationType = this.OperationType,
+                RuCharges = ruCharges,
+                CosmosDiagnostics = readResponse.Diagnostics,
+                LazyDiagnostics = () => $"read {readResponse.Diagnostics} write: {writeDiagnostics}",
+            };
+        }
+
+        public Task PrepareAsync()
+        {
+            string newPartitionKey = Guid.NewGuid().ToString();
+            this.sampleJObject["id"] = Guid.NewGuid().ToString();
+            this.sampleJObject[this.partitionKeyPath] = newPartitionKey;
+            return Task.CompletedTask;
+        }
+    }
+}