fix(ENGKNOW-3778): fix sporadic CRAM reference failures and hangs

model/src/main/java/org/gorpipe/gor/driver/providers/stream/datatypes/cram/CramIterator.java

@@ -278,12 +278,14 @@ private CRAMReferenceSource createFileReference(File refFile, boolean preferFold
         if (refFile.isDirectory()) {
             return createCompositeReferenceSource(refFile);
         } else if (preferFolder) {
-            try {
-                return createCompositeReferenceSource(refFile.getParentFile());
-            } catch (Exception e) {
-                // Fallback to single file, in case none of the files contains proper meta.
-                return createSharedFastaReferenceSource(refFile);
+            var folderSource = new FolderReferenceSource(refFile.getParentFile().getPath());
+            if (folderSource.containsReferenceFile(refFile.toPath())) {
+                return createCompositeReferenceSource(folderSource, refFile.getParentFile());
             }
+            // Fallback to single file, in case the reference file does not contain proper meta.
+            log.debug("Reference file {} is not md5 indexed, using it as single fasta reference", refFile.getPath());
+            folderSource.close();
+            return createSharedFastaReferenceSource(refFile);
         } else {
             return createSharedFastaReferenceSource(refFile);
         }
@@ -298,9 +300,13 @@ private CRAMReferenceSource createSharedFastaReferenceSource(File refFile) {
         return new SharedFastaReferenceSource(referenceFile, referenceKey);
     }
     private CRAMReferenceSource createCompositeReferenceSource(File refFolder) {
+        return createCompositeReferenceSource(new FolderReferenceSource(refFolder.getPath()), refFolder);
+    }
+
+    private CRAMReferenceSource createCompositeReferenceSource(FolderReferenceSource folderSource, File refFolder) {
         log.debug("Using folder reference for CRAM: {}", refFolder.getPath());
         return new CompositeReferenceSource(List.of(
-                new FolderReferenceSource(refFolder.getPath()),
+                folderSource,
                 new EBIReferenceSource(refFolder.getPath())));
     }
 }

model/src/main/java/org/gorpipe/gor/driver/providers/stream/datatypes/cram/reference/EBIReferenceSource.java

@@ -14,8 +14,10 @@
 import java.io.IOException;
 import java.io.InputStream;
 import java.net.URL;
+import java.net.URLConnection;
 import java.nio.file.Files;
 import java.nio.file.Path;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.Locale;
 import java.util.Map;
@@ -30,15 +32,19 @@ public class EBIReferenceSource extends MD5CachedReferenceSource {
     private static final Logger log = LoggerFactory.getLogger(EBIReferenceSource.class);
 
     public static final String KEY_USE_CRAM_REF_DOWNLOAD = "gor.driver.cram.ref.download";
+    public static final String KEY_CRAM_REF_DOWNLOAD_URL_MASK = "gor.driver.cram.ref.download.urlmask";
+    public static final String KEY_CRAM_REF_DOWNLOAD_TIMEOUT = "gor.driver.cram.ref.download.timeout"; // Milliseconds
 
     private static final String REFBASES_PREFIX = "md5_";
     private static final String REFBASES_EXT = ".txt";
 
-    protected static Map<String, Path> md5ToRefbases = new ConcurrentHashMap<>();
+    // Shared between instances, as scanning the folder is expensive.
+    private static final Map<Path, Md5FolderIndex<Path>> indexByFolder = new ConcurrentHashMap<>();
 
     private static final int DOWNLOAD_TRIES_BEFORE_FAILING = 2;
 
     private Path referenceFolder;  // If null we do not download.
+    private Md5FolderIndex<Path> md5ToRefbases;  // Null if no reference folder.
 
     public EBIReferenceSource() {
     }
@@ -50,45 +56,37 @@ public EBIReferenceSource(String referenceFolder) {
             if (!Files.isDirectory(this.referenceFolder)) {
                 throw new GorResourceException("Can not create FolderReferenceSource %s as the target is not a folder or does not exists".formatted(referenceFolder), referenceFolder);
             }
-            scanReferenceFolder();
+            md5ToRefbases = Md5FolderIndex.get(indexByFolder, this.referenceFolder, EBIReferenceSource::scanReferenceFolder);
+            md5ToRefbases.refreshIfStale();
         }
     }
 
     Set<Path> getRefbasesFiles() {
-        return new HashSet<>(md5ToRefbases.values());
+        return md5ToRefbases != null ? new HashSet<>(md5ToRefbases.values()) : new HashSet<>();
     }
 
-    private void scanReferenceFolder() {
-        md5ToRefbases.clear();
-
-        if (referenceFolder == null) return;
-
-        try {
-            for (var p : Files.list(referenceFolder).filter(Files::isRegularFile).toList()) {
+    private static Map<String, Path> scanReferenceFolder(Path referenceFolder) {
+        Map<String, Path> md5Map = new HashMap<>();
+        try (var files = Files.list(referenceFolder)) {
+            for (var p : files.filter(Files::isRegularFile).toList()) {
                 var f = p.getFileName().toString().toLowerCase();
                 if (f.startsWith(REFBASES_PREFIX) && f.endsWith(REFBASES_EXT)) {
-                    processRefbasesFile(referenceFolder.resolve(f));
+                    String md5 = f.substring(REFBASES_PREFIX.length(), f.length() - REFBASES_EXT.length());
+                    md5Map.put(md5, p);
                 }
             }
         } catch (IOException e) {
             log.warn("Failed scanning reference folder {}", referenceFolder, e);
         }
-    }
-
-    private void processRefbasesFile(Path refbases) {
-        String fileName = refbases.getFileName().toString();
-        if (!fileName.startsWith(REFBASES_PREFIX) || !fileName.endsWith(REFBASES_EXT)) return;
-
-        String md5 = fileName.substring(REFBASES_PREFIX.length(), fileName.length() - REFBASES_EXT.length());
-        md5ToRefbases.put(md5, refbases);
+        return md5Map;
     }
 
     @Override
     protected byte[] loadReference(final SAMSequenceRecord record) {
         var md5 = record.getMd5();
 
         // Load from refbases file.
-        Path refbasesPath = md5ToRefbases.get(md5);
+        Path refbasesPath = md5ToRefbases != null ? md5ToRefbases.get(md5) : null;
         if (refbasesPath != null) {
             try {
                 byte[] bases = Files.readAllBytes(refbasesPath);
@@ -125,10 +123,16 @@ protected byte[] loadReference(final SAMSequenceRecord record) {
      * @throws IOException    if the sequence is not found or the download fails.
      */
     private byte[] downloadFromEBI(final String md5) throws IOException {
-        final String url = String.format(Locale.US, Defaults.EBI_REFERENCE_SERVICE_URL_MASK, md5);
+        final String urlMask = System.getProperty(KEY_CRAM_REF_DOWNLOAD_URL_MASK, Defaults.EBI_REFERENCE_SERVICE_URL_MASK);
+        final String url = String.format(Locale.US, urlMask, md5);
+        // Without timeouts an unresponsive server blocks forever (while holding the reference source lock).
+        final int timeout = Integer.parseInt(System.getProperty(KEY_CRAM_REF_DOWNLOAD_TIMEOUT, "30000"));
 
         for (int i = 0; i < DOWNLOAD_TRIES_BEFORE_FAILING; i++) {
-            try (final InputStream is = new URL(url).openStream()) {
+            final URLConnection connection = new URL(url).openConnection();
+            connection.setConnectTimeout(timeout);
+            connection.setReadTimeout(timeout);
+            try (final InputStream is = connection.getInputStream()) {
                 if (is == null)
                     return null;
 

model/src/main/java/org/gorpipe/gor/driver/providers/stream/datatypes/cram/reference/FolderReferenceSource.java

@@ -49,9 +49,12 @@ public class FolderReferenceSource extends MD5CachedReferenceSource {
     private static final Set<String> FASTA_EXT = Set.of("fa", "fasta");
 
     private record Md5ReferencePath(String md5, Path path, String contig) {}
-    protected static Map<String, Md5ReferencePath> md5ToReferencePath = new ConcurrentHashMap<>();
 
-    private Path referenceFolder;
+    // Shared between instances, as scanning the folder is expensive.
+    private static final Map<Path, Md5FolderIndex<Md5ReferencePath>> indexByFolder = new ConcurrentHashMap<>();
+
+    private final Path referenceFolder;
+    private final Md5FolderIndex<Md5ReferencePath> md5ToReferencePath;
 
     private final Map<Path, ReferenceSequenceFile> refFileByPath = new ConcurrentHashMap<>();
 
@@ -60,7 +63,8 @@ public FolderReferenceSource(String referenceFolder) {
         if (!Files.isDirectory(this.referenceFolder)) {
             throw new GorResourceException("Can not create FolderReferenceSource %s as the target is not a folder or does not exists".formatted(referenceFolder), referenceFolder);
         }
-        scanReferenceFolder();
+        md5ToReferencePath = Md5FolderIndex.get(indexByFolder, this.referenceFolder, FolderReferenceSource::scanReferenceFolder);
+        md5ToReferencePath.refreshIfStale();
     }
 
     @Override
@@ -80,31 +84,37 @@ protected byte[] loadReference(final SAMSequenceRecord record) {
         return null;
     }
 
-    private void scanReferenceFolder() {
-        md5ToReferencePath.clear();
-
-        try {
-            for (var p : Files.list(referenceFolder).filter(Files::isRegularFile).toList()) {
+    private static Map<String, Md5ReferencePath> scanReferenceFolder(Path referenceFolder) {
+        Map<String, Md5ReferencePath> md5Map = new HashMap<>();
+        try (var files = Files.list(referenceFolder)) {
+            for (var p : files.filter(Files::isRegularFile).toList()) {
                 var f = p.getFileName().toString().toLowerCase();
                 if (FASTA_EXT.stream().anyMatch(ext -> f.endsWith("." + ext))) {
-                    processFasta(referenceFolder.resolve(f));
+                    // One bad fasta file should not prevent using the others.
+                    try {
+                        processFasta(p, md5Map);
+                    } catch (Exception e) {
+                        log.error("Skipping invalid cram reference file {}", p, e);
+                    }
                 }
             }
         } catch (IOException e) {
             log.warn("Failed scanning reference folder {}", referenceFolder, e);
         }
+        return md5Map;
     }
 
-    private void processFasta(Path fastaFile) {
-        ReferenceSequenceFile refFile = refFileByPath.computeIfAbsent(fastaFile, ReferenceSequenceFileFactory::getReferenceSequenceFile);
-        SAMSequenceDictionary dictionary = refFile.getSequenceDictionary();
-        if (dictionary == null) {
-            throw new GorResourceException("Fasta file %s is invalid cram reference as it is missing dict file".formatted(fastaFile), fastaFile.toString());
-        }
-        for (SAMSequenceRecord rec : dictionary.getSequences()) {
-            String md5 = rec.getMd5();
-            if (md5 == null || md5.isEmpty()) continue;
-            md5ToReferencePath.put(md5, new FolderReferenceSource.Md5ReferencePath(md5, fastaFile, rec.getContig()));
+    private static void processFasta(Path fastaFile, Map<String, Md5ReferencePath> md5Map) throws IOException {
+        try (ReferenceSequenceFile refFile = ReferenceSequenceFileFactory.getReferenceSequenceFile(fastaFile)) {
+            SAMSequenceDictionary dictionary = refFile.getSequenceDictionary();
+            if (dictionary == null) {
+                throw new GorResourceException("Fasta file %s is invalid cram reference as it is missing dict file".formatted(fastaFile), fastaFile.toString());
+            }
+            for (SAMSequenceRecord rec : dictionary.getSequences()) {
+                String md5 = rec.getMd5();
+                if (md5 == null || md5.isEmpty()) continue;
+                md5Map.put(md5, new FolderReferenceSource.Md5ReferencePath(md5, fastaFile, rec.getContig()));
+            }
         }
     }
 
@@ -119,7 +129,17 @@ public void close() {
         refFileByPath.clear();
     }
 
+    /**
+     * @return true if the given fasta file is a valid (md5 indexed) reference file in this folder.
+     */
+    public boolean containsReferenceFile(Path fastaFile) {
+        Path normalized = fastaFile.toAbsolutePath().normalize();
+        return getReferenceFiles().stream().anyMatch(p -> p.toAbsolutePath().normalize().equals(normalized));
+    }
+
     Set<Path> getReferenceFiles() {
-        return new HashSet<>(refFileByPath.keySet());
+        Set<Path> files = new HashSet<>();
+        md5ToReferencePath.values().forEach(r -> files.add(r.path()));
+        return files;
     }
 }

model/src/main/java/org/gorpipe/gor/driver/providers/stream/datatypes/cram/reference/Md5FolderIndex.java

@@ -0,0 +1,65 @@
+package org.gorpipe.gor.driver.providers.stream.datatypes.cram.reference;
+
+import java.nio.file.Path;
+import java.util.Collection;
+import java.util.Map;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.Function;
+
+/**
+ * Index from md5 to reference data for a reference folder.  Shared by all reference sources using the same folder,
+ * as scanning the folder is expensive.
+ *
+ * The folder is rescanned when the index is older than the rescan interval.  A rescan builds a new map and swaps it in,
+ * so readers never see a partially built (or cleared) index (see ENGKNOW-3778).
+ */
+class Md5FolderIndex<V> {
+
+    public static final String KEY_RESCAN_INTERVAL = "gor.driver.cram.reference.rescaninterval"; // Seconds
+
+    private static final long NOT_SCANNED = -1;
+
+    private final Path folder;
+    private final Function<Path, Map<String, V>> scanner;
+
+    private volatile Map<String, V> md5Map = new ConcurrentHashMap<>();
+    private long lastScanMillis = NOT_SCANNED;
+
+    /**
+     * Get the index for the given folder, creating it if needed.
+     */
+    static <V> Md5FolderIndex<V> get(Map<Path, Md5FolderIndex<V>> indexByFolder, Path folder,
+                                     Function<Path, Map<String, V>> scanner) {
+        return indexByFolder.computeIfAbsent(folder.toAbsolutePath().normalize(),
+                f -> new Md5FolderIndex<>(f, scanner));
+    }
+
+    private Md5FolderIndex(Path folder, Function<Path, Map<String, V>> scanner) {
+        this.folder = folder;
+        this.scanner = scanner;
+    }
+
+    V get(String md5) {
+        return md5 != null ? md5Map.get(md5) : null;
+    }
+
+    synchronized void put(String md5, V value) {
+        md5Map.put(md5, value);
+    }
+
+    Collection<V> values() {
+        return md5Map.values();
+    }
+
+    /**
+     * Rescan the folder if it has not been scanned or the index is older than the rescan interval.
+     */
+    synchronized void refreshIfStale() {
+        long now = System.currentTimeMillis();
+        long intervalMillis = Long.parseLong(System.getProperty(KEY_RESCAN_INTERVAL, "300")) * 1000;
+        if (lastScanMillis == NOT_SCANNED || now - lastScanMillis >= intervalMillis) {
+            md5Map = new ConcurrentHashMap<>(scanner.apply(folder));
+            lastScanMillis = now;
+        }
+    }
+}