Support user-defined protocol for benchmarks.

atomix · Dec 13, 2018 · 94f3764 · 94f3764
1 parent bc405c1
commit 94f3764
Show file tree

Hide file tree

Showing 2 changed files with 25 additions and 13 deletions.
diff --git a/bench/src/main/java/io/atomix/bench/BenchmarkConfig.java b/bench/src/main/java/io/atomix/bench/BenchmarkConfig.java
@@ -16,6 +16,7 @@
 package io.atomix.bench;
 
 import io.atomix.primitive.config.PrimitiveConfig;
+import io.atomix.primitive.protocol.PrimitiveProtocolConfig;
 
 import java.util.UUID;
 
@@ -33,7 +34,7 @@ public class BenchmarkConfig {
   private static final boolean DEFAULT_DETERMINISTIC = true;
 
   private String benchId = UUID.randomUUID().toString();
-  private PrimitiveConfig primitive;
+  private PrimitiveProtocolConfig protocol;
   private int operations = DEFAULT_OPERATIONS;
   private int writePercentage = DEFAULT_WRITE_PERCENTAGE;
   private int numKeys = DEFAULT_NUM_KEYS;
@@ -52,12 +53,12 @@ public BenchmarkConfig setBenchId(String benchId) {
     return this;
   }
 
-  public PrimitiveConfig getPrimitive() {
+  public PrimitiveProtocolConfig getProtocol() {
-    return primitive;
+    return protocol;
   }
 
-  public BenchmarkConfig setPrimitive(PrimitiveConfig primitive) {
+  public BenchmarkConfig setProtocol(PrimitiveProtocolConfig protocol) {
-    this.primitive = primitive;
+    this.protocol = protocol;
     return this;
   }
 

diff --git a/bench/src/main/java/io/atomix/bench/BenchmarkRunner.java b/bench/src/main/java/io/atomix/bench/BenchmarkRunner.java
@@ -17,6 +17,8 @@
 
 import io.atomix.core.Atomix;
 import io.atomix.core.map.AtomicMap;
+import io.atomix.primitive.protocol.PrimitiveProtocol;
+import io.atomix.primitive.protocol.ProxyProtocol;
 import io.atomix.utils.serializer.Namespaces;
 import io.atomix.utils.serializer.Serializer;
 import org.slf4j.Logger;
@@ -162,9 +164,22 @@ public void run() {
       teardown();
     }
 
+    void read(String key) {
+      map.get(key);
+      readCounter.incrementAndGet();
+    }
+
+    void write(String key, String value) {
+      map.put(key, value);
+      writeCounter.incrementAndGet();
+    }
+
+    @SuppressWarnings("unchecked")
     void setup() {
+      ProxyProtocol protocol = (ProxyProtocol) atomix.getRegistry().getType(PrimitiveProtocol.Type.class, "foo").newProtocol(config.getProtocol());
       map = atomix.<String, String>atomicMapBuilder("bench")
           .withSerializer(Serializer.using(Namespaces.BASIC))
+          .withProtocol(protocol)
           .build();
       if (config.isIncludeEvents()) {
         map.addListener(event -> {
@@ -184,11 +199,9 @@ private class NonDeterministicSubmitter extends Submitter {
     void submit() {
       String key = keys[random.nextInt(keys.length)];
       if (random.nextInt(100) < config.getWritePercentage()) {
-        map.put(key, values[random.nextInt(values.length)]);
+        write(key, values[random.nextInt(values.length)]);
-        writeCounter.incrementAndGet();
       } else {
-        map.get(key);
+        read(key);
-        readCounter.incrementAndGet();
       }
     }
   }
@@ -199,11 +212,9 @@ private class DeterministicSubmitter extends Submitter {
     @Override
     void submit() {
       if (random.nextInt(100) < config.getWritePercentage()) {
-        map.put(keys[index++ % keys.length], values[random.nextInt(values.length)]);
+        write(keys[index++ % keys.length], values[random.nextInt(values.length)]);
-        writeCounter.incrementAndGet();
       } else {
-        map.get(keys[random.nextInt(keys.length)]);
+        read(keys[random.nextInt(keys.length)]);
-        readCounter.incrementAndGet();
       }
     }
   }