X-Git-Url: https://git.arvados.org/arvados.git/blobdiff_plain/85bd5f3272cfc8ba8e54f4f47ea9865d7de62779..8de691c25eac0454f8f30cfa35eccff15642e330:/sdk/cwl/arvados_cwl/fsaccess.py

diff --git a/sdk/cwl/arvados_cwl/fsaccess.py b/sdk/cwl/arvados_cwl/fsaccess.py
index 93e2819084..5e3d685911 100644
--- a/sdk/cwl/arvados_cwl/fsaccess.py
+++ b/sdk/cwl/arvados_cwl/fsaccess.py
@@ -9,6 +9,7 @@ import urlparse
 import re
 import logging
 import threading
+from collections import OrderedDict
 
 import ruamel.yaml as yaml
 
@@ -26,19 +27,44 @@ from schema_salad.ref_resolver import DefaultFetcher
 logger = logging.getLogger('arvados.cwl-runner')
 
 class CollectionCache(object):
-    def __init__(self, api_client, keep_client, num_retries):
+    def __init__(self, api_client, keep_client, num_retries,
+                 cap=256*1024*1024,
+                 min_entries=2):
         self.api_client = api_client
         self.keep_client = keep_client
-        self.collections = {}
+        self.collections = OrderedDict()
         self.lock = threading.Lock()
+        self.total = 0
+        self.cap = cap
+        self.min_entries = min_entries
+
+    def cap_cache(self):
+        if self.total > self.cap:
+            # ordered list iterates from oldest to newest
+            for pdh, v in self.collections.items():
+                if self.total < self.cap or len(self.collections) < self.min_entries:
+                    break
+                # cut it loose
+                logger.debug("Evicting collection reader %s from cache", pdh)
+                del self.collections[pdh]
+                self.total -= v[1]
 
     def get(self, pdh):
         with self.lock:
             if pdh not in self.collections:
                 logger.debug("Creating collection reader for %s", pdh)
-                self.collections[pdh] = arvados.collection.CollectionReader(pdh, api_client=self.api_client,
-                                                                            keep_client=self.keep_client)
-            return self.collections[pdh]
+                cr = arvados.collection.CollectionReader(pdh, api_client=self.api_client,
+                                                         keep_client=self.keep_client)
+                sz = len(cr.manifest_text()) * 128
+                self.collections[pdh] = (cr, sz)
+                self.total += sz
+                self.cap_cache()
+            else:
+                cr, sz = self.collections[pdh]
+                # bump it to the back
+                del self.collections[pdh]
+                self.collections[pdh] = (cr, sz)
+            return cr
 
 
 class CollectionFsAccess(cwltool.stdfsaccess.StdFsAccess):
@@ -81,14 +107,14 @@ class CollectionFsAccess(cwltool.stdfsaccess.StdFsAccess):
 
     def glob(self, pattern):
         collection, rest = self.get_collection(pattern)
-        if collection and not rest:
+        if collection is not None and not rest:
             return [pattern]
         patternsegments = rest.split("/")
         return self._match(collection, patternsegments, "keep:" + collection.manifest_locator())
 
     def open(self, fn, mode):
         collection, rest = self.get_collection(fn)
-        if collection:
+        if collection is not None:
             return collection.open(rest, mode)
         else:
             return super(CollectionFsAccess, self).open(self._abs(fn), mode)
@@ -105,7 +131,7 @@ class CollectionFsAccess(cwltool.stdfsaccess.StdFsAccess):
 
     def isfile(self, fn):  # type: (unicode) -> bool
         collection, rest = self.get_collection(fn)
-        if collection:
+        if collection is not None:
             if rest:
                 return isinstance(collection.find(rest), arvados.arvfile.ArvadosFile)
             else:
@@ -115,7 +141,7 @@ class CollectionFsAccess(cwltool.stdfsaccess.StdFsAccess):
 
     def isdir(self, fn):  # type: (unicode) -> bool
         collection, rest = self.get_collection(fn)
-        if collection:
+        if collection is not None:
             if rest:
                 return isinstance(collection.find(rest), arvados.collection.RichCollectionBase)
             else:
@@ -125,7 +151,7 @@ class CollectionFsAccess(cwltool.stdfsaccess.StdFsAccess):
 
     def listdir(self, fn):  # type: (unicode) -> List[unicode]
         collection, rest = self.get_collection(fn)
-        if collection:
+        if collection is not None:
             if rest:
                 dir = collection.find(rest)
             else:
@@ -147,7 +173,7 @@ class CollectionFsAccess(cwltool.stdfsaccess.StdFsAccess):
         if path.startswith("$(task.tmpdir)") or path.startswith("$(task.outdir)"):
             return path
         collection, rest = self.get_collection(path)
-        if collection:
+        if collection is not None:
             return path
         else:
             return os.path.realpath(path)
@@ -227,6 +253,9 @@ workflow_uuid_pattern = re.compile(r'[a-z0-9]{5}-7fd4e-[a-z0-9]{15}')
 pipeline_template_uuid_pattern = re.compile(r'[a-z0-9]{5}-p5p6p-[a-z0-9]{15}')
 
 def collectionResolver(api_client, document_loader, uri, num_retries=4):
+    if uri.startswith("keep:") or uri.startswith("arvwf:"):
+        return uri
+
     if workflow_uuid_pattern.match(uri):
         return "arvwf:%s#main" % (uri)