Merge pull request #50 from 0aix/parallel

Add option to run queries in parallel

Author: 0aix

README.md

@@ -183,6 +183,12 @@ You may want queries, asserts or query metadata to use a specific date column, o
 
 Simply pass `--parameter KEY=VALUE` in the CLI and the `KEY` will be replaced with `VALUE` in all queries, query metadata and test assertions. For example, to query June 23rd 2015, you could use `--parameter DATE=2015-06-23`. If the query uses `${DATE}` in the query it will be replaced before execution with `2015-06-23`.
 
+### Query Parallelism
+
+You may want to run queries in parallel rather than sequentially especially if you have many time-consuming queries.
+
+To enable this feature, pass in `--query-parallel-enable true` when launching Validatar. By default, this will run all queries in parallel. If this number needs to be limited, pass in `--query-parallel-max VALUE` where `VALUE` is the max number of queries that should run concurrently.  
+
 ## Execution Engines
 
 ### Hive
@@ -330,6 +336,17 @@ Option                                 Description
   fully qualified classes to plug in>
 
 
+Engine Options:
+Option                              Description                           
+------                              -----------                           
+--query-parallel-enable <Boolean:   Whether or not queries should run in  
+  Query parallelism option>           parallel. (default: false)          
+--query-parallel-max <Integer: Max  The max number of queries that will   
+  query parallelism>                  run concurrently. If non-positive or
+                                      unspecified, all queries will run at
+                                      once. (default: 0)                  
+
+
 Hive engine options:
 Option (* = required)                   Description
 ---------------------                   -----------
@@ -488,10 +505,11 @@ Version | Notes
 0.5.4   | Added a flag ```--report-on-failure-only``` to only generate reports if there were failures in tests (including warnOnly) or queries
 0.5.5   | Shaded ```org.objectweb.asm``` to not clash with asm in Hadoop environments
 0.5.6   | Fixed a bug with pretty-printing results with nulls
-0.6.0   | Better reporting (show data with only the assertion columns with the assertion result column, columns now in sorted order, EMail Subject Prefix). Can now write multiple reports per invocation (specify more than one report formatter using --report-format)
+0.6.0   | Better reporting (show data with only the assertion columns with the assertion result column, columns now in sorted order, Email Subject Prefix). Can now write multiple reports per invocation (specify more than one report formatter using --report-format)
 0.6.1   | Added a flag to configure the Email reporter SMTP strategy. Use ```--email-smtp-strategy``` to pass in ```SMTP_PLAIN```, ```SMTP_TLS``` or ```SMTP_SSL```.
 0.6.2   | Bintray EOL. First rerelease of 0.6.1 on Maven Central instead
 0.6.3   | Screwdriver migration. First rerelease of 0.6.1 using Screwdriver instead of Travis.
+0.6.4   | Added support for running queries in parallel
 
 ## Members
 

pom.xml

@@ -401,7 +401,7 @@
       <dependency>
           <groupId>org.apache.httpcomponents</groupId>
           <artifactId>httpclient</artifactId>
-          <version>4.5.1</version>
+          <version>4.5.13</version>
       </dependency>
       <dependency>
           <groupId>com.google.guava</groupId>

src/main/java/com/yahoo/validatar/execution/EngineManager.java

@@ -11,13 +11,16 @@
 import com.yahoo.validatar.execution.hive.Apiary;
 import com.yahoo.validatar.execution.pig.Sty;
 import com.yahoo.validatar.execution.rest.JSON;
+import joptsimple.OptionParser;
+import joptsimple.OptionSet;
 import lombok.extern.slf4j.Slf4j;
 
 import java.util.Arrays;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
+import java.util.concurrent.ForkJoinPool;
 import java.util.stream.Collectors;
 
 /**
@@ -27,6 +30,29 @@
 public class EngineManager extends Pluggable<Engine> implements Helpable {
     public static final String CUSTOM_ENGINE = "custom-engine";
     public static final String CUSTOM_ENGINE_DESCRIPTION = "Additional custom engine to load.";
+    public static final String QUERY_PARALLEL_ENABLE = "query-parallel-enable";
+    public static final String QUERY_PARALLEL_MAX = "query-parallel-max";
+    private static final int QUERY_PARALLEL_MIN = 1;
+
+    protected boolean queryParallelEnable;
+    protected int queryParallelMax;
+
+    private static final OptionParser PARSER = new OptionParser() {
+        {
+            accepts(QUERY_PARALLEL_ENABLE, "Whether or not queries should run in parallel.")
+                    .withRequiredArg()
+                    .describedAs("Query parallelism option")
+                    .ofType(Boolean.class)
+                    .defaultsTo(false);
+            accepts(QUERY_PARALLEL_MAX, "The max number of queries that will run concurrently. If non-positive or " +
+                                        "unspecified, all queries will run at once.")
+                    .withRequiredArg()
+                    .describedAs("Max query parallelism")
+                    .ofType(Integer.class)
+                    .defaultsTo(0);
+            allowsUnrecognizedOptions();
+        }
+    };
 
     /**
      * The Engine classes to manage.
@@ -85,6 +111,10 @@ public EngineManager(String[] arguments) {
             engines.put(engine.getName(), new WorkingEngine(engine));
             log.info("Added engine {} to list of engines.", engine.getName());
         }
+
+        OptionSet parser = PARSER.parse(arguments);
+        queryParallelEnable = (Boolean) parser.valueOf(QUERY_PARALLEL_ENABLE);
+        queryParallelMax = (Integer) parser.valueOf(QUERY_PARALLEL_MAX);
     }
 
     /**
@@ -107,8 +137,7 @@ protected boolean startEngines(List<Query> queries) {
         List<Query> all = queries == null ? Collections.emptyList() : queries;
         // Queries -> engine name Set -> start engine -> verify all started
         return all.stream().map(q -> q.engine).collect(Collectors.toSet())
-               .stream().map(this::startEngine)
-               .allMatch(b -> b);
+               .stream().allMatch(this::startEngine);
     }
 
     private boolean startEngine(String engine) {
@@ -132,6 +161,7 @@ private boolean startEngine(String engine) {
 
     @Override
     public void printHelp() {
+        Helpable.printHelp("Engine Options", PARSER);
         engines.values().stream().map(WorkingEngine::getEngine).forEach(Engine::printHelp);
         Helpable.printHelp("Advanced Engine Options", getPluginOptionsParser());
     }
@@ -155,7 +185,19 @@ public boolean run(List<Query> queries) {
             return false;
         }
         // Run each query.
-        queries.stream().forEach(this::run);
+        if (!queryParallelEnable) {
+            queries.forEach(this::run);
+        } else {
+            int poolSize = Math.max(queryParallelMax > 0 ? queryParallelMax : queries.size(), QUERY_PARALLEL_MIN);
+            log.info("Creating a ForkJoinPool with size {}", poolSize);
+            ForkJoinPool forkJoinPool = new ForkJoinPool(poolSize);
+            try {
+                forkJoinPool.submit(() -> queries.parallelStream().forEach(this::run)).get();
+            } catch (Exception e) {
+                log.error("Caught exception", e);
+            }
+            forkJoinPool.shutdown();
+        }
         return true;
     }
 }

src/main/java/com/yahoo/validatar/execution/rest/JSON.java

@@ -35,6 +35,7 @@
 import java.util.List;
 import java.util.Map;
 import java.util.Objects;
+import java.util.concurrent.TimeUnit;
 
 import static java.util.Arrays.asList;
 
@@ -283,9 +284,11 @@ private HttpClient createClient(Map<String, String> metadata) {
         int retries = Integer.valueOf(metadata.getOrDefault(METADATA_RETRY_KEY, String.valueOf(defaultRetries)));
         RequestConfig config = RequestConfig.custom().setConnectTimeout(timeout)
                                                      .setConnectionRequestTimeout(timeout)
-                                                     .setSocketTimeout(timeout).build();
+                                                     .setSocketTimeout(timeout)
+                                                     .build();
         return HttpClientBuilder.create()
                                 .setDefaultRequestConfig(config)
+                                .setConnectionTimeToLive(timeout, TimeUnit.MILLISECONDS)
                                 .setRetryHandler(new DefaultHttpRequestRetryHandler(retries, false))
                                 .build();
     }

src/test/java/com/yahoo/validatar/execution/EngineManagerTest.java

@@ -255,4 +255,45 @@ public void testNormalRun() {
         Assert.assertEquals((String) actual.get("Foo.b").get(0).data, (String) expected.get("Foo.b").get(0).data);
         Assert.assertEquals((String) actual.get("Foo.b").get(1).data, (String) expected.get("Foo.b").get(1).data);
     }
+
+    @Test
+    public void testParallelRun() {
+        query.engine = MockRunningEngine.ENGINE_NAME;
+        query.name = "Foo";
+
+        String[] args = {"--query-parallel-enable", "true"};
+        manager = new EngineManager(args);
+        manager.setEngines(engines);
+
+        Assert.assertTrue(manager.run(queries));
+
+        Map<String, List<TypedObject>> expected = new HashMap<>();
+        List<TypedObject> columns = new ArrayList<>();
+        columns.add(new TypedObject("42", TypeSystem.Type.STRING));
+        expected.put("Foo.a", columns);
+        columns = new ArrayList<>();
+        columns.add(new TypedObject("42", TypeSystem.Type.STRING));
+        columns.add(new TypedObject("52", TypeSystem.Type.STRING));
+        expected.put("Foo.b", columns);
+
+        Map<String, Column> actual = query.getResult().getColumns();
+
+        Assert.assertEquals(actual.size(), 2);
+        Assert.assertEquals(expected.size(), 2);
+        Assert.assertEquals(actual.get("Foo.a").size(), 1);
+        Assert.assertEquals(expected.get("Foo.a").size(), 1);
+        Assert.assertEquals((String) actual.get("Foo.a").get(0).data, (String) expected.get("Foo.a").get(0).data);
+        Assert.assertEquals(actual.get("Foo.b").size(), 2);
+        Assert.assertEquals(expected.get("Foo.b").size(), 2);
+        Assert.assertEquals((String) actual.get("Foo.b").get(0).data, (String) expected.get("Foo.b").get(0).data);
+        Assert.assertEquals((String) actual.get("Foo.b").get(1).data, (String) expected.get("Foo.b").get(1).data);
+    }
+
+    @Test
+    public void testConstructor() {
+        String[] args = {"--query-parallel-enable", "true", "--query-parallel-max", "10"};
+        manager = new EngineManager(args);
+        Assert.assertTrue(manager.queryParallelEnable);
+        Assert.assertEquals(manager.queryParallelMax, 10);
+    }
 }