broadinstitute · andersleung · Jun 16, 2020 · Jun 16, 2020 · Jun 16, 2020 · Jun 17, 2020
diff --git a/src/main/java/org/broadinstitute/hellbender/engine/GATKPath.java b/src/main/java/org/broadinstitute/hellbender/engine/GATKPath.java
@@ -26,6 +26,7 @@ public class GATKPath extends PathSpecifier implements TaggedArgument, Serializa
     private static final long serialVersionUID = 1L;
 
     public static final String HDFS_SCHEME = "hdfs";
+    public static final String HTSGET_SCHEME = "htsget";
 
     private String tagName;
     private Map<String, String> tagAttributes;

diff --git a/src/main/java/org/broadinstitute/hellbender/engine/GATKTool.java b/src/main/java/org/broadinstitute/hellbender/engine/GATKTool.java
@@ -13,6 +13,7 @@
 import java.nio.file.Path;
 import java.time.ZonedDateTime;
 import java.util.*;
+import java.util.stream.Collectors;
 import java.util.stream.Stream;
 import org.broadinstitute.barclay.argparser.Argument;
 import org.broadinstitute.barclay.argparser.ArgumentCollection;
@@ -455,8 +456,16 @@ else if (hasCramInput()) {
                 factory = factory.enable(SamReaderFactory.Option.CACHE_FILE_BASED_INDEXES);
             }
 
-            reads = new ReadsPathDataSource(readArguments.getReadPaths(), readArguments.getReadIndexPaths(), factory, cloudPrefetchBuffer,
-                (cloudIndexPrefetchBuffer < 0 ? cloudPrefetchBuffer : cloudIndexPrefetchBuffer));
+            final Map<String, List<GATKPath>> pathsByScheme = readArguments.getReadPathSpecifiers().stream()
+                .collect(Collectors.groupingBy(path -> path.getURI().getScheme()));
+
+            if (pathsByScheme.get(GATKPath.HTSGET_SCHEME) != null) {
+                reads = new ReadsHtsgetDataSource(pathsByScheme.get(GATKPath.HTSGET_SCHEME), factory);
+            } else {
+                reads = new ReadsPathDataSource(readArguments.getReadPaths(), readArguments.getReadIndexPaths(), factory, cloudPrefetchBuffer,
+                    (cloudIndexPrefetchBuffer < 0 ? cloudPrefetchBuffer : cloudIndexPrefetchBuffer));
+            }
+
         }
         else {
             reads = null;

diff --git a/src/main/java/org/broadinstitute/hellbender/engine/ReadsHtsgetDataSource.java b/src/main/java/org/broadinstitute/hellbender/engine/ReadsHtsgetDataSource.java
diff --git a/src/main/java/org/broadinstitute/hellbender/engine/filters/IntervalOverlapReadFilter.java b/src/main/java/org/broadinstitute/hellbender/engine/filters/IntervalOverlapReadFilter.java
@@ -1,15 +1,12 @@
 package org.broadinstitute.hellbender.engine.filters;
 
+import htsjdk.samtools.util.Locatable;
 import htsjdk.samtools.util.OverlapDetector;
 import org.broadinstitute.barclay.argparser.Advanced;
 import org.broadinstitute.barclay.argparser.Argument;
 import org.broadinstitute.barclay.help.DocumentedFeature;
 import org.broadinstitute.hellbender.cmdline.ReadFilterArgumentDefinitions;
-import org.broadinstitute.hellbender.utils.GenomeLocParser;
-import org.broadinstitute.hellbender.utils.GenomeLocSortedSet;
-import org.broadinstitute.hellbender.utils.IntervalMergingRule;
-import org.broadinstitute.hellbender.utils.IntervalSetRule;
-import org.broadinstitute.hellbender.utils.IntervalUtils;
+import org.broadinstitute.hellbender.utils.*;
 import org.broadinstitute.hellbender.utils.help.HelpConstants;
 import org.broadinstitute.hellbender.utils.logging.OneShotLogger;
 import org.broadinstitute.hellbender.utils.read.GATKRead;

diff --git a/src/main/java/org/broadinstitute/hellbender/tools/HtsgetReader.java b/src/main/java/org/broadinstitute/hellbender/tools/HtsgetReader.java
@@ -6,29 +6,12 @@
 import java.io.InputStream;
 import java.io.OutputStream;
 import java.net.URI;
-import java.nio.charset.Charset;
-import java.nio.charset.StandardCharsets;
-import java.util.ArrayList;
+import java.nio.file.Files;
+import java.nio.file.Path;
 import java.util.List;
-import java.util.concurrent.ExecutionException;
-import java.util.concurrent.ExecutorService;
-import java.util.concurrent.Executors;
-import java.util.concurrent.Future;
-import java.util.concurrent.ThreadFactory;
+import java.util.stream.Stream;
 
-import com.fasterxml.jackson.databind.DeserializationFeature;
-import com.fasterxml.jackson.databind.MapperFeature;
-import com.fasterxml.jackson.databind.ObjectMapper;
-import com.google.common.util.concurrent.ThreadFactoryBuilder;
-
-import org.apache.commons.io.Charsets;
 import org.apache.commons.io.IOUtils;
-import org.apache.http.Header;
-import org.apache.http.HttpEntity;
-import org.apache.http.client.methods.CloseableHttpResponse;
-import org.apache.http.client.methods.HttpGet;
-import org.apache.http.impl.client.CloseableHttpClient;
-import org.apache.http.util.EntityUtils;
 import org.broadinstitute.barclay.argparser.Advanced;
 import org.broadinstitute.barclay.argparser.Argument;
 import org.broadinstitute.barclay.argparser.CommandLineProgramProperties;
@@ -37,13 +20,7 @@
 import org.broadinstitute.hellbender.cmdline.StandardArgumentDefinitions;
 import org.broadinstitute.hellbender.cmdline.programgroups.ExampleProgramGroup;
 import org.broadinstitute.hellbender.exceptions.UserException;
-import org.broadinstitute.hellbender.tools.htsgetreader.HtsgetClass;
-import org.broadinstitute.hellbender.tools.htsgetreader.HtsgetErrorResponse;
-import org.broadinstitute.hellbender.tools.htsgetreader.HtsgetFormat;
-import org.broadinstitute.hellbender.tools.htsgetreader.HtsgetRequestBuilder;
-import org.broadinstitute.hellbender.tools.htsgetreader.HtsgetRequestField;
-import org.broadinstitute.hellbender.tools.htsgetreader.HtsgetResponse;
-import org.broadinstitute.hellbender.utils.HttpUtils;
+import org.broadinstitute.hellbender.tools.htsgetreader.*;
 import org.broadinstitute.hellbender.utils.SimpleInterval;
 import org.broadinstitute.hellbender.utils.Utils;
 
@@ -75,7 +52,7 @@ public class HtsgetReader extends CommandLineProgram {
     public static final String FIELDS_LONG_NAME = "field";
     public static final String TAGS_LONG_NAME = "tag";
     public static final String NOTAGS_LONG_NAME = "notag";
-    public static final String NUM_THREADS_LONG_NAME = "reader-threads";
+    public static final String PARALLEL_DOWNLOAD_LONG_NAME = "parallel";
     public static final String CHECK_MD5_LONG_NAME = "check-md5";
 
     @Argument(doc = "Output file.",
@@ -130,161 +107,74 @@ public class HtsgetReader extends CommandLineProgram {
     private List<String> notags;
 
     @Advanced
-    @Argument(fullName = NUM_THREADS_LONG_NAME,
-            shortName = NUM_THREADS_LONG_NAME,
-            doc = "How many simultaneous threads to use when reading data from an htsget response;" +
-                    "higher values may improve performance when network latency is an issue.",
-            optional = true,
-            minValue = 1)
-    private int readerThreads = 1;
+    @Argument(doc = "Whether to try to download blocks in parallel",
+        fullName = PARALLEL_DOWNLOAD_LONG_NAME,
+        shortName = PARALLEL_DOWNLOAD_LONG_NAME,
+        optional = true)
+    private final boolean parallelDownload = false;
 
     @Argument(fullName = CHECK_MD5_LONG_NAME, shortName = CHECK_MD5_LONG_NAME, doc = "Boolean determining whether to calculate the md5 digest of the assembled file "
             + "and validate it against the provided md5 hash, if it exists.", optional = true)
     private boolean checkMd5 = false;
 
-    private ExecutorService executorService;
-
-    private CloseableHttpClient client;
-
-    @Override
-    public void onStartup() {
-        if (this.readerThreads > 1) {
-            logger.info("Initializing with " + this.readerThreads + " threads");
-            final ThreadFactory threadFactory = new ThreadFactoryBuilder()
-                .setNameFormat("htsgetReader-thread-%d")
-                .setDaemon(true).build();
-            this.executorService = Executors.newFixedThreadPool(readerThreads, threadFactory);
-        }
-        this.client = HttpUtils.getClient();
-    }
-
-    @Override
-    public void onShutdown() {
-        if (this.executorService != null) {
-            this.executorService.shutdownNow();
-        }
-        super.onShutdown();
-    }
-
-    /**
-     * Downloads data blocks provided by response to outputFile in serial
-     */
-    private void getData(final HtsgetResponse response) {
-        try (final OutputStream ostream = new FileOutputStream(this.outputFile)) {
-            response.getBlocks().forEach(b -> {
-                try (final InputStream istream = b.getData()) {
-                    IOUtils.copy(istream, ostream);
-                } catch (final IOException e) {
-                    throw new UserException("Failed to copy data block to output file", e);
-                }
-            });
-        } catch (final IOException e) {
-            throw new UserException("Could not create output file: " + outputFile, e);
-        }
-    }
-
-    /**
-     * Downloads data blocks provided by response to outputFile in parallel, using
-     * the number of threads specified by user
-     */
-    private void getDataParallel(final HtsgetResponse response) {
-        final List<Future<InputStream>> futures = new ArrayList<>(response.getBlocks().size());
-        response.getBlocks().forEach(b -> futures.add(this.executorService.submit(b::getData)));
-
-        try (final OutputStream ostream = new FileOutputStream(this.outputFile)) {
-            futures.forEach(f -> {
-                try (final InputStream istream = f.get()) {
-                    IOUtils.copy(istream, ostream);
-                } catch (final IOException e) {
-                    throw new UserException("Error while copying data block to output file", e);
-                } catch (final ExecutionException | InterruptedException e) {
-                    throw new UserException("Error while waiting to download block", e);
-                }
-            });
-        } catch (final IOException e) {
-            throw new UserException("Could not create output file", e);
-        }
-    }
-
-    /**
-     * Checks md5 digest provided in response, if one exists, against calculated md5
-     * hash of downloaded file, warning user if they differ
-     */
-    private void checkMd5(final HtsgetResponse resp) {
-        final String expectedMd5 = resp.getMd5();
-        if (expectedMd5 == null) {
-            logger.warn("No md5 digest provided by response");
-        } else {
-            try {
-                final String actualMd5 = Utils.calculateFileMD5(outputFile);
-                if (!actualMd5.equals(expectedMd5)) {
-                    throw new UserException("Expected md5: " + expectedMd5 + " did not match actual md5: " + actualMd5);
-                }
-            } catch (final IOException e) {
-                throw new UserException("Unable to calculate md5 digest", e);
+    private void checkMd5(final String expectedMd5) {
+        try {
+            final String actualMd5 = Utils.calculateFileMD5(this.outputFile);
+            if (!actualMd5.equals(expectedMd5)) {
+                throw new UserException("Expected md5: " + expectedMd5 + " and actual md5: " + actualMd5 + " do not match");
             }
+        } catch (final IOException e) {
+            throw new UserException("Could not calculate md5 checksum from downloaded file", e);
         }
     }
 
-    private ObjectMapper getObjectMapper() {
-        final ObjectMapper mapper = new ObjectMapper();
-        mapper.enable(DeserializationFeature.UNWRAP_ROOT_VALUE);
-        mapper.configure(MapperFeature.ACCEPT_CASE_INSENSITIVE_PROPERTIES, true);
-        return mapper;
-    }
-
     @Override
     public Object doWork() {
-        // construct request from command line args and convert to URI
-        final HtsgetRequestBuilder req = new HtsgetRequestBuilder(endpoint, id)
+        // Construct request from command line args and convert to URI
+        final HtsgetRequest req = new HtsgetRequest(endpoint, id)
             .withFormat(format)
             .withDataClass(dataClass)
             .withInterval(interval)
             .withFields(fields)
             .withTags(tags)
             .withNotags(notags);
-        final URI reqURI = req.toURI();
-
-        final HttpGet getReq = new HttpGet(reqURI);
-        try (final CloseableHttpResponse resp = this.client.execute(getReq)) {
-            // get content of response
-            final HttpEntity entity = resp.getEntity();
-            final Header encodingHeader = entity.getContentEncoding();
-            final Charset encoding = encodingHeader == null 
-                ? StandardCharsets.UTF_8
-                : Charsets.toCharset(encodingHeader.getValue());
-            final String jsonBody = EntityUtils.toString(entity, encoding);
-
-            final ObjectMapper mapper = this.getObjectMapper();
 
-            if (resp.getStatusLine() == null) {
-                throw new UserException("htsget server response did not contain status line");
-            }
-            final int statusCode = resp.getStatusLine().getStatusCode();
-            if (400 <= statusCode && statusCode < 500) {
-                final HtsgetErrorResponse err = mapper.readValue(jsonBody, HtsgetErrorResponse.class);
-                throw new UserException("Invalid request, received error code: " + statusCode + ", error type: "
-                        + err.getError() + ", message: " + err.getMessage());
-            } else if (statusCode == 200) {
-                final HtsgetResponse response = mapper.readValue(jsonBody, HtsgetResponse.class);
-
-                if (this.readerThreads > 1) {
-                    this.getDataParallel(response);
-                } else {
-                    this.getData(response);
-                }
-
-                logger.info("Successfully wrote to: " + outputFile);
+        final HtsgetResponse resp = req.getResponse();
+        if (resp.getMd5() == null) {
+            this.checkMd5 = false;
+            logger.info("No md5 checksum received");
+        }
 
-                if (checkMd5) {
-                    this.checkMd5(response);
-                }
+        try (final OutputStream outputstream = new FileOutputStream(this.outputFile)) {
+            if (this.parallelDownload) {
+                resp.getBlocks().parallelStream().map(block -> {
+                        final Path tempFile = org.broadinstitute.hellbender.utils.io.IOUtils.createTempPath("htsget-temp", "");
+                        try (final OutputStream ostream = Files.newOutputStream(tempFile)) {
+                            org.apache.commons.io.IOUtils.copy(block.getData(), ostream);
+                            return Files.newInputStream(tempFile);
+                        } catch (final IOException e) {
+                            throw new UserException("Error while downloading htsget block", e);
+                        }
+                    }).forEachOrdered(inputStream -> {
+                        try {
+                            IOUtils.copy(inputStream, outputstream);
+                        } catch (final IOException e) {
+                            throw new UserException("IOException while writing output file", e);
+                        }
+                    });
             } else {
-                throw new UserException("Unrecognized status code: " + statusCode);
+                try {
+                    IOUtils.copy(resp.getDataStream(), outputstream);
+                } catch (final IOException e) {
+                    throw new UserException("IOException while writing output file", e);
+                }
             }
         } catch (final IOException e) {
             throw new UserException("IOException during htsget download", e);
         }
+
+        if (this.checkMd5) this.checkMd5(resp.getMd5());
+
         return null;
     }
 }