Reworked 3.x async version.

Both code and commit history will be cleaned up later.

Author: Bouncheck

benchmarks/basic/java-driver-3.x-async/source/src/main/java/Config.java

@@ -14,7 +14,6 @@ enum Workload {
     long tasks;
     long concurrency;
     boolean dont_prepare;
-    int threads;
 
     Config(String[] args)  {
 
@@ -23,16 +22,14 @@ enum Workload {
         this.tasks = 1000 * 1000;
         this.concurrency = 1024;
         this.dont_prepare = false;
-        this.threads = Runtime.getRuntime().availableProcessors();
 
         Options options = new Options();
 
         options.addOption("d", "dont-prepare", false, "Don't create tables and insert into them before the benchmark");
         options.addOption("n", "nodes", true, "Addresses of database nodes to connect to separated by a comma");
         options.addOption("w", "workload", true, "Type of work to perform (Inserts, Selects, Mixed)");
         options.addOption("t", "tasks", true, "Total number of tasks (requests) to perform the during benchmark. In case of mixed workload there will be tasks inserts and tasks selects");
-        options.addOption("c", "concurrency", true, "Maximum number of requests performed at once (concurrency/threads has to fit in java int)");
-        options.addOption("th", "threads", true, "Number of worker threads to launch. If omitted defaults to number of cores JVM can see");
+        options.addOption("c", "concurrency", true, "Maximum number of requests performed at once");
 
         try {
             CommandLineParser parser = new DefaultParser();
@@ -60,9 +57,6 @@ enum Workload {
                 this.concurrency = Integer.parseInt(cmd.getOptionValue("concurrency"));
             }
 
-            if(cmd.hasOption("threads")){
-                this.threads = Integer.parseInt(cmd.getOptionValue("threads"));
-            }
         } catch (ParseException e){
             HelpFormatter helpFormatter = new HelpFormatter();
             helpFormatter.printHelp("./run.sh [OPTION]...", options);
@@ -79,7 +73,6 @@ public String toString() {
                 ", tasks=" + tasks +
                 ", concurrency=" + concurrency +
                 ", dont_prepare=" + dont_prepare +
-                ", threads=" + threads +
                 '}';
     }
 }

benchmarks/basic/java-driver-3.x-async/source/src/main/java/MainClass.java

@@ -1,13 +1,11 @@
 import com.datastax.driver.core.*;
-import com.google.common.base.Function;
 import com.google.common.util.concurrent.AsyncFunction;
 import com.google.common.util.concurrent.FutureCallback;
 import com.google.common.util.concurrent.Futures;
 import com.google.common.util.concurrent.ListenableFuture;
 
 import java.util.ArrayList;
 import java.util.concurrent.*;
-import java.util.concurrent.atomic.AtomicLong;
 
 public class MainClass {
 
@@ -18,6 +16,9 @@ public class MainClass {
     private static final String INSERT_STRING = "INSERT INTO benchks.benchtab (pk, v1, v2) VALUES(?, ?, ?)";
     private static final String SELECT_STRING = "SELECT v1, v2 FROM benchks.benchtab WHERE pk = ?";
 
+    private static PreparedStatement INSERT_PS;
+    private static PreparedStatement SELECT_PS;
+
     private static long benchmarkStart;
 
     private static long benchmarkEnd;
@@ -42,87 +43,34 @@ public static void main(String[] args) throws InterruptedException, ExecutionExc
             }
         }
 
-        AtomicLong requestsSent = new AtomicLong(0);
 
-        executor = Executors.newFixedThreadPool(config.threads);
+        ArrayList<CompletableFuture<ResultSet>> arr = new ArrayList<>();
 
         System.out.println("Starting the benchmark");
 
         benchmarkStart = System.nanoTime();
-        ArrayList<Future<?>> arr = new ArrayList<>();
-        for (int i = 0; i < config.threads; i++) {
-            arr.add(
-                    executor.submit(() -> {
-                        int permits = (int) (config.concurrency / config.threads);
-                        Semaphore semaphore = new Semaphore(permits);
-                        PreparedStatement insertQ = session.prepare(INSERT_STRING);
-                        PreparedStatement selectQ = session.prepare(SELECT_STRING);
-                        long pk;
-                        while ((pk = requestsSent.incrementAndGet()) < config.tasks) {
-                            try {
-                                semaphore.acquire();
-                            } catch (InterruptedException e) {
-                                throw new RuntimeException(e);
-                            }
-                            if (config.workload.equals(Config.Workload.Inserts) || config.workload.equals(Config.Workload.Mixed)) {
-                                ListenableFuture<ResultSet> resultSetA = session.executeAsync(insertQ.bind(pk, 2L * pk, 3L * pk));
-                                ListenableFuture<ResultSet> resultSetB;
-
-                                long finalPk = pk;
-                                if (config.workload.equals(Config.Workload.Selects) || config.workload.equals(Config.Workload.Mixed)) {
-                                    // Chain with select depending on workload
-                                    resultSetB = Futures.transform(resultSetA,
-                                            new AsyncFunction<ResultSet, ResultSet>() {
-                                                public ListenableFuture<ResultSet> apply(ResultSet rs) throws Exception {
-                                                    return session.executeAsync(selectQ.bind(finalPk));
-                                                }
-                                            });
-                                } else {
-                                    resultSetB = resultSetA;
-                                }
-                                ListenableFuture<Boolean> result;
-                                // Verify returned data depending on workload and return the permit
-                                if (config.workload.equals(Config.Workload.Selects) || config.workload.equals(Config.Workload.Mixed)) {
-                                    result = Futures.transform(resultSetB,
-                                            new Function<ResultSet, Boolean>() {
-                                                public Boolean apply(ResultSet rs) {
-                                                    semaphore.release();
-                                                    Row r = rs.one();
-                                                    if ((r.getLong("v1") != 2 * finalPk) || (r.getLong("v2") != 3 * finalPk)) {
-                                                        throw new RuntimeException(String.format("Received incorrect data. " + "Expected: (%s, %s, %s). " + "Received: (%s, %s ,%s).", finalPk, 2 * finalPk, 3 * finalPk, r.getInt("pk"), r.getInt("v1"), r.getInt("v2")));
-                                                    }
-                                                    return true;
-                                                }
-                                            });
-
-                                }
-                                else {
-                                    result = Futures.transform(resultSetB,
-                                            new Function<ResultSet, Boolean>() {
-                                                public Boolean apply(ResultSet rs) {
-                                                    semaphore.release();
-                                                    return true;
-                                                }
-                                            });
-                                }
-                            }
-                        }
-                        try {
-                            //possible only if all futures ended successfully
-                            semaphore.acquire(permits);
-                        } catch (InterruptedException e) {
-                            throw new RuntimeException(e);
-                        }
-                    }));
-        }
 
+        INSERT_PS = session.prepare(INSERT_STRING);
+        SELECT_PS = session.prepare(SELECT_STRING);
+
+        if (config.workload == Config.Workload.Inserts) {
+            for (int i = 0; i < config.concurrency; i++) {
+                arr.add(executeInsert(session, i * (config.tasks / config.concurrency), (i + 1) * (config.tasks / config.concurrency)));
+            }
+        } else if (config.workload == Config.Workload.Selects) {
+            for (int i = 0; i < config.concurrency; i++) {
+                arr.add(executeSelect(session, i * (config.tasks / config.concurrency), (i + 1) * (config.tasks / config.concurrency)));
+            }
+        } else if (config.workload == Config.Workload.Mixed) {
+            for (int i = 0; i < config.concurrency; i++) {
+                arr.add(executeMixed(session, i * (config.tasks / config.concurrency), (i + 1) * (config.tasks / config.concurrency)));
+            }
+        }
 
-        executor.shutdown();
-        executor.awaitTermination(Long.MAX_VALUE, TimeUnit.NANOSECONDS);
-        benchmarkEnd = System.nanoTime();
         for (Future<?> f : arr) {
-            f.get(); // make sure nothing has thrown
+            f.get(); // make sure nothing has thrown and everything finished
         }
+        benchmarkEnd = System.nanoTime();
         System.out.println(String.format("Finished\nBenchmark time: %d ms\n", (benchmarkEnd - benchmarkStart) / 1_000_000));
 
         session.close();
@@ -141,33 +89,99 @@ static void prepareKeyspaceAndTable(Session session) {
     private static void prepareSelectsBenchmark(Session session) throws InterruptedException, ExecutionException {
         System.out.println("Preparing a selects benchmark (inserting values)...");
 
-        AtomicLong nextBatchStart = new AtomicLong(0);
-        executor = Executors.newFixedThreadPool((int) config.threads);
+        ArrayList<CompletableFuture<ResultSet>> arr = new ArrayList<>();
+        INSERT_PS = session.prepare(INSERT_STRING);
 
-        ArrayList<Future<?>> arr = new ArrayList<>();
-        try {
-            for (int i = 0; i < config.concurrency; i++) {
-                arr.add(executor.submit(() -> {
-                    PreparedStatement insertQ = session.prepare(INSERT_STRING);
-                    while (true) {
-                        long curBatchStart = nextBatchStart.addAndGet(config.tasks / config.concurrency);
-                        if (curBatchStart >= config.tasks) {
-                            break;
-                        }
-                        long curBatchEnd = Math.min(curBatchStart + (config.tasks / config.concurrency), config.tasks);
-                        for (long pk = curBatchStart; pk < curBatchEnd; pk++) {
-                            session.execute(insertQ.bind(pk, 2L * pk, 3L * pk));
-                        }
-                    }
-                }));
+        for (int i = 0; i < config.concurrency; i++) {
+            arr.add(executeInsert(session, i * (config.tasks / config.concurrency), (i + 1) * (config.tasks / config.concurrency)));
+        }
+        for (Future<?> f : arr) {
+            f.get(); // make sure nothing has thrown and everything finished
+        }
+
+    }
+
+    public static CompletableFuture<ResultSet> executeSelect(Session s, long currentIter, long maxIter) {
+        if (currentIter >= maxIter) {
+            // No more iterations
+            return CompletableFuture.completedFuture(null);
+        }
+
+        ResultSetFuture fut = s.executeAsync(SELECT_PS.bind(currentIter));
+
+        // Convert ResultSetFuture to CompletableFuture
+        CompletableFuture<ResultSet> futCompletable = new CompletableFuture<>();
+        Futures.addCallback(fut, new FutureCallback<ResultSet>() {
+            @Override
+            public void onSuccess(ResultSet result) {
+                futCompletable.complete(result);
+            }
+
+            @Override
+            public void onFailure(Throwable t) {
+                futCompletable.completeExceptionally(t);
+            }
+        });
+
+        // Execute next iteration after that
+        return futCompletable.thenCompose(rs -> executeSelect(s, currentIter + 1, maxIter));
+    }
+
+    public static CompletableFuture<ResultSet> executeInsert(Session s, long currentIter, long maxIter) {
+        if (currentIter >= maxIter) {
+            // No more iterations
+            return CompletableFuture.completedFuture(null);
+        }
+
+        ResultSetFuture fut = s.executeAsync(INSERT_PS.bind(currentIter, 2L * currentIter, 3L * currentIter));
+
+        // Convert ResultSetFuture to CompletableFuture
+        CompletableFuture<ResultSet> futCompletable = new CompletableFuture<>();
+        Futures.addCallback(fut, new FutureCallback<ResultSet>() {
+            @Override
+            public void onSuccess(ResultSet result) {
+                futCompletable.complete(result);
             }
-        } finally {
-            executor.shutdown();
-            executor.awaitTermination(Long.MAX_VALUE, TimeUnit.NANOSECONDS);
-            for (Future<?> f : arr) {
-                f.get(); // make sure nothing has thrown
+
+            @Override
+            public void onFailure(Throwable t) {
+                futCompletable.completeExceptionally(t);
             }
+        });
+
+        // Execute next iteration after that
+        return futCompletable.thenCompose(rs -> executeInsert(s, currentIter + 1, maxIter));
+    }
+
+    public static CompletableFuture<ResultSet> executeMixed(Session s, long currentIter, long maxIter) {
+        if (currentIter >= maxIter) {
+            // No more iterations
+            return CompletableFuture.completedFuture(null);
         }
+
+        ListenableFuture<ResultSet> listenableFut = s.executeAsync(INSERT_PS.bind(currentIter, 2L * currentIter, 3L * currentIter));
+        listenableFut = Futures.transform(listenableFut, new AsyncFunction<ResultSet, ResultSet>() {
+            public ListenableFuture<ResultSet> apply(ResultSet rs) throws Exception {
+                return (s.executeAsync(SELECT_PS.bind(currentIter)));
+            }
+        });
+
+        // Convert ResultSetFuture to CompletableFuture
+        CompletableFuture<ResultSet> futCompletable = new CompletableFuture<>();
+        Futures.addCallback(listenableFut, new FutureCallback<ResultSet>() {
+            @Override
+            public void onSuccess(ResultSet result) {
+                futCompletable.complete(result);
+            }
+
+            @Override
+            public void onFailure(Throwable t) {
+                futCompletable.completeExceptionally(t);
+            }
+        });
+
+        // Execute next iteration after that
+        return futCompletable.thenCompose(rs -> executeMixed(s, currentIter + 1, maxIter));
     }
 }