Adding cosmos encryption benchmarking

sdk/cosmos/azure-cosmos-benchmark/pom.xml

@@ -46,6 +46,28 @@ Licensed under the MIT License.
       <version>4.18.0-beta.1</version> <!-- {x-version-update;com.azure:azure-cosmos;current} -->
     </dependency>
 
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-cosmos-encryption</artifactId>
+      <version>1.0.0-beta.7</version> <!-- {x-version-update;com.azure:azure-cosmos-encryption;current} -->
+    </dependency>
+
+    <dependency>
+      <groupId>com.azure</groupId>
+      <artifactId>azure-identity</artifactId>
+      <version>1.3.4</version>  <!-- {x-version-update;com.azure:azure-identity;dependency} -->
+      <exclusions>
+        <exclusion>
+          <groupId>com.azure</groupId>
+          <artifactId>azure-core</artifactId>
+        </exclusion>
+        <exclusion>
+          <groupId>com.azure</groupId>
+          <artifactId>azure-core-http-netty</artifactId>
+        </exclusion>
+      </exclusions>
+    </dependency>
+
     <dependency>
       <groupId>com.beust</groupId>
       <artifactId>jcommander</artifactId>

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/AsyncQueryBenchmark.java

@@ -82,7 +82,7 @@ protected void performWorkload(BaseSubscriber<FeedResponse<PojoizedJson>> baseSu
         } else if (configuration.getOperationType() == Configuration.Operation.QuerySingle) {
 
             int index = r.nextInt(this.configuration.getNumberOfPreCreatedDocuments());
-            String pk = docsToRead.get(index).getProperty(partitionKey);
+            String pk = (String) docsToRead.get(index).getProperty(partitionKey);
             options.setPartitionKey(new PartitionKey(pk));
             String sqlQuery = "Select * from c where c." + partitionKey + " = \"" + pk + "\"";
             obs = cosmosAsyncContainer.queryItems(sqlQuery, options, PojoizedJson.class).byPage();
@@ -126,7 +126,7 @@ protected void performWorkload(BaseSubscriber<FeedResponse<PojoizedJson>> baseSu
         } else if (configuration.getOperationType() == Configuration.Operation.ReadAllItemsOfLogicalPartition) {
 
             int index = r.nextInt(this.configuration.getNumberOfPreCreatedDocuments());
-            String pk = docsToRead.get(index).getProperty(partitionKey);
+            String pk = (String) docsToRead.get(index).getProperty(partitionKey);
             obs = cosmosAsyncContainer.readAllItems(new PartitionKey(pk), options, PojoizedJson.class).byPage();
         } else {
             throw new IllegalArgumentException("Unsupported Operation: " + configuration.getOperationType());

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/BenchmarkHelper.java

@@ -10,7 +10,7 @@ public class BenchmarkHelper {
     public static PojoizedJson generateDocument(String idString, String dataFieldValue, String partitionKey,
                                          int dataFieldCount) {
         PojoizedJson instance = new PojoizedJson();
-        Map<String, String> properties = instance.getInstance();
+        Map<String, Object> properties = instance.getInstance();
         properties.put("id", idString);
         properties.put(partitionKey, idString);
         for (int i = 0; i < dataFieldCount; i++) {

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/Configuration.java

@@ -96,6 +96,18 @@ public class Configuration {
     @Parameter(names = "-preferredRegionsList", description = "Comma separated preferred regions list")
     private String preferredRegionsList;
 
+    @Parameter(names = "-encryptedStringFieldCount", description = "Number of string field that need to be encrypted")
+    private int encryptedStringFieldCount = 1;
+
+    @Parameter(names = "-encryptedLongFieldCount", description = "Number of long field that need to be encrypted")
+    private int encryptedLongFieldCount = 0;
+
+    @Parameter(names = "-encryptedDoubleFieldCount", description = "Number of double field that need to be encrypted")
+    private int encryptedDoubleFieldCount = 0;
+
+    @Parameter(names = "-encryptionEnabled", description = "Control switch to enable the encryption operation")
+    private boolean encryptionEnabled = false;
+
     @Parameter(names = "-operation", description = "Type of Workload:\n"
         + "\tReadThroughput- run a READ workload that prints only throughput *\n"
         + "\tReadThroughputWithMultipleClients - run a READ workload that prints throughput and latency for multiple client read.*\n"
@@ -461,6 +473,22 @@ public List<String> getPreferredRegionsList() {
         return preferredRegions;
     }
 
+    public int getEncryptedStringFieldCount() {
+        return encryptedStringFieldCount;
+    }
+
+    public int getEncryptedLongFieldCount() {
+        return encryptedLongFieldCount;
+    }
+
+    public int getEncryptedDoubleFieldCount() {
+        return encryptedDoubleFieldCount;
+    }
+
+    public boolean isEncryptionEnabled() {
+        return encryptionEnabled;
+    }
+
     public void tryGetValuesFromSystem() {
         serviceEndpoint = StringUtils.defaultString(Strings.emptyToNull(System.getenv().get("SERVICE_END_POINT")),
                                                     serviceEndpoint);
@@ -500,7 +528,24 @@ public void tryGetValuesFromSystem() {
                 Strings.emptyToNull(System.getenv().get("THROUGHPUT")), Integer.toString(throughput));
         throughput = Integer.parseInt(throughputValue);
 
-        preferredRegionsList = StringUtils.defaultString(Strings.emptyToNull(System.getenv().get("PREFERRED_REGIONS_LIST")), preferredRegionsList);
+        preferredRegionsList = StringUtils.defaultString(Strings.emptyToNull(System.getenv().get(
+            "PREFERRED_REGIONS_LIST")), preferredRegionsList);
+
+        encryptedStringFieldCount = Integer.parseInt(
+            StringUtils.defaultString(Strings.emptyToNull(System.getenv().get("ENCRYPTED_STRING_FIELD_COUNT")),
+                Integer.toString(encryptedStringFieldCount)));
+
+        encryptedLongFieldCount = Integer.parseInt(
+            StringUtils.defaultString(Strings.emptyToNull(System.getenv().get("ENCRYPTED_LONG_FIELD_COUNT")),
+                Integer.toString(encryptedLongFieldCount)));
+
+        encryptedDoubleFieldCount = Integer.parseInt(
+            StringUtils.defaultString(Strings.emptyToNull(System.getenv().get("ENCRYPTED_DOUBLE_FIELD_COUNT")),
+                Integer.toString(encryptedDoubleFieldCount)));
+
+        encryptionEnabled = Boolean.parseBoolean(StringUtils.defaultString(Strings.emptyToNull(System.getenv().get(
+            "ENCRYPTED_ENABLED")),
+            Boolean.toString(encryptionEnabled)));
     }
 
     private synchronized MeterRegistry azureMonitorMeterRegistry(String instrumentationKey) {

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/Main.java

@@ -4,13 +4,19 @@
 package com.azure.cosmos.benchmark;
 
 import com.azure.cosmos.benchmark.ctl.AsyncCtlWorkload;
+import com.azure.cosmos.benchmark.encryption.AsyncEncryptionBenchmark;
+import com.azure.cosmos.benchmark.encryption.AsyncEncryptionQueryBenchmark;
+import com.azure.cosmos.benchmark.encryption.AsyncEncryptionQuerySinglePartitionMultiple;
+import com.azure.cosmos.benchmark.encryption.AsyncEncryptionReadBenchmark;
+import com.azure.cosmos.benchmark.encryption.AsyncEncryptionWriteBenchmark;
 import com.azure.cosmos.benchmark.linkedin.LICtlWorkload;
 import com.beust.jcommander.JCommander;
 import com.beust.jcommander.ParameterException;
-import java.util.Optional;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import java.util.Optional;
+
 import static com.azure.cosmos.benchmark.Configuration.Operation.CtlWorkload;
 import static com.azure.cosmos.benchmark.Configuration.Operation.LinkedInCtlWorkload;
 import static com.azure.cosmos.benchmark.Configuration.Operation.ReadThroughputWithMultipleClients;
@@ -37,14 +43,15 @@ public static void main(String[] args) throws Exception {
             if (cfg.isSync()) {
                 syncBenchmark(cfg);
             } else {
-                if(cfg.getOperationType().equals(ReadThroughputWithMultipleClients)) {
+                if (cfg.getOperationType().equals(ReadThroughputWithMultipleClients)) {
                     asyncMultiClientBenchmark(cfg);
-                } else if(cfg.getOperationType().equals(CtlWorkload)) {
+                } else if (cfg.getOperationType().equals(CtlWorkload)) {
                     asyncCtlWorkload(cfg);
                 } else if (cfg.getOperationType().equals(LinkedInCtlWorkload)) {
                     linkedInCtlWorkload(cfg);
-                }
-                else {
+                } else if (cfg.isEncryptionEnabled()) {
+                    asyncEncryptionBenchmark(cfg);
+                } else {
                     asyncBenchmark(cfg);
                 }
             }
@@ -159,6 +166,47 @@ private static void asyncBenchmark(Configuration cfg) throws Exception {
         }
     }
 
+    private static void asyncEncryptionBenchmark(Configuration cfg) throws Exception {
+        LOGGER.info("Async encryption benchmark ...");
+        AsyncEncryptionBenchmark<?> benchmark = null;
+        try {
+            switch (cfg.getOperationType()) {
+                case WriteThroughput:
+                case WriteLatency:
+                    benchmark = new AsyncEncryptionWriteBenchmark(cfg);
+                    break;
+
+                case ReadThroughput:
+                case ReadLatency:
+                    benchmark = new AsyncEncryptionReadBenchmark(cfg);
+                    break;
+
+                case QueryCross:
+                case QuerySingle:
+                case QueryParallel:
+                case QueryOrderby:
+                case QueryTopOrderby:
+                case QueryInClauseParallel:
+                    benchmark = new AsyncEncryptionQueryBenchmark(cfg);
+                    break;
+
+                case QuerySingleMany:
+                    benchmark = new AsyncEncryptionQuerySinglePartitionMultiple(cfg);
+                    break;
+
+                default:
+                    throw new RuntimeException(cfg.getOperationType() + " is not supported");
+            }
+
+            LOGGER.info("Starting {}", cfg.getOperationType());
+            benchmark.run();
+        } finally {
+            if (benchmark != null) {
+                benchmark.shutdown();
+            }
+        }
+    }
+
     private static void asyncMultiClientBenchmark(Configuration cfg) throws Exception {
         LOGGER.info("Async multi client benchmark ...");
         AsynReadWithMultipleClients<?> benchmark = null;

sdk/cosmos/azure-cosmos-benchmark/src/main/java/com/azure/cosmos/benchmark/PojoizedJson.java

@@ -7,30 +7,31 @@
 import com.fasterxml.jackson.annotation.JsonAnyGetter;
 import com.fasterxml.jackson.annotation.JsonAnySetter;
 import com.fasterxml.jackson.annotation.JsonIgnore;
+import com.fasterxml.jackson.annotation.JsonProperty;
 
 import java.util.HashMap;
 import java.util.Map;
 
 public class PojoizedJson {
-    private final Map<String, String> instanceProps = new HashMap<>();
+    private final Map<String, Object> instanceProps = new HashMap<>();
 
     @JsonAnyGetter
-    public Map<String, String> getInstance() {
+    public Map<String, Object> getInstance() {
         return instanceProps;
     }
 
     @JsonAnySetter
-    public void setProperty(String name, String value) {
+    public void setProperty(String name, Object value) {
         this.instanceProps.put(name, value);
     }
 
     @JsonIgnore
     public String getId() {
-        return instanceProps.get("id");
+        return (String) instanceProps.get("id");
     }
 
     @JsonIgnore
-    public String getProperty(String propName) {
+    public Object getProperty(String propName) {
         return instanceProps.get(propName);
     }
 }