2803: Split arvados and arvados_fuse Python modules.
[arvados.git] / sdk / python / arvados / keep.py
index 5f9a6d916b7839f17e1d3e76a9dd43b3a4c5b1fe..88487ae96e672726cfa5dbb3142dcfdeafcbec94 100644 (file)
@@ -17,6 +17,7 @@ import zlib
 import fcntl
 import time
 import threading
+import timer
 
 global_client_object = None
 
@@ -141,20 +142,23 @@ class KeepClient(object):
         self._cache_lock = threading.Lock()
         self._cache = []
         # default 256 megabyte cache
-        self._cache_max = 256 * 1024 * 1024
+        self.cache_max = 256 * 1024 * 1024
 
     def shuffled_service_roots(self, hash):
         if self.service_roots == None:
             self.lock.acquire()
-            keep_disks = api().keep_disks().list().execute()['items']
-            roots = (("http%s://%s:%d/" %
-                      ('s' if f['service_ssl_flag'] else '',
-                       f['service_host'],
-                       f['service_port']))
-                     for f in keep_disks)
-            self.service_roots = sorted(set(roots))
-            logging.debug(str(self.service_roots))
-            self.lock.release()
+            try:
+                keep_disks = arvados.api().keep_disks().list().execute()['items']
+                roots = (("http%s://%s:%d/" %
+                          ('s' if f['service_ssl_flag'] else '',
+                           f['service_host'],
+                           f['service_port']))
+                         for f in keep_disks)
+                self.service_roots = sorted(set(roots))
+                logging.debug(str(self.service_roots))
+            finally:
+                self.lock.release()
+
         seed = hash
         pool = self.service_roots[:]
         pseq = []
@@ -171,41 +175,115 @@ class KeepClient(object):
         logging.debug(str(pseq))
         return pseq
 
+    class CacheSlot(object):
+        def __init__(self, locator):
+            self.locator = locator
+            self.ready = threading.Event()
+            self.content = None
+
+        def get(self):
+            self.ready.wait()
+            return self.content
+
+        def set(self, value):
+            self.content = value
+            self.ready.set()
+
+        def size(self):
+            if self.content == None:
+                return 0
+            else:
+                return len(self.content)
+
+    def cap_cache(self):
+        '''Cap the cache size to self.cache_max'''
+        self._cache_lock.acquire()
+        try:
+            self._cache = filter(lambda c: not (c.ready.is_set() and c.content == None), self._cache)
+            sm = sum([slot.size() for slot in self._cache])
+            while sm > self.cache_max:
+                del self._cache[-1]
+                sm = sum([slot.size() for a in self._cache])
+        finally:
+            self._cache_lock.release()
+
+    def reserve_cache(self, locator):
+        '''Reserve a cache slot for the specified locator, 
+        or return the existing slot.'''
+        self._cache_lock.acquire()
+        try:
+            # Test if the locator is already in the cache
+            for i in xrange(0, len(self._cache)):
+                if self._cache[i].locator == locator:
+                    n = self._cache[i]
+                    if i != 0:
+                        # move it to the front
+                        del self._cache[i]
+                        self._cache.insert(0, n)
+                    return n, False
+
+            # Add a new cache slot for the locator
+            n = KeepClient.CacheSlot(locator)
+            self._cache.insert(0, n)
+            return n, True
+        finally:
+            self._cache_lock.release()
+
     def get(self, locator):
+        #logging.debug("Keep.get %s" % (locator))
+
         if re.search(r',', locator):
             return ''.join(self.get(x) for x in locator.split(','))
         if 'KEEP_LOCAL_STORE' in os.environ:
             return KeepClient.local_store_get(locator)
         expect_hash = re.sub(r'\+.*', '', locator)
 
-        c = self.check_cache(expect_hash)
-        if c:
-            return c
-
-        for service_root in self.shuffled_service_roots(expect_hash):
-            url = service_root + expect_hash
-            api_token = config.get('ARVADOS_API_TOKEN')
-            headers = {'Authorization': "OAuth2 %s" % api_token,
-                       'Accept': 'application/octet-stream'}
-            blob = self.get_url(url, headers, expect_hash)
-            if blob:
-                self.put_cache(expect_hash, blob)
-                return blob
-
-        for location_hint in re.finditer(r'\+K@([a-z0-9]+)', locator):
-            instance = location_hint.group(1)
-            url = 'http://keep.' + instance + '.arvadosapi.com/' + expect_hash
-            blob = self.get_url(url, {}, expect_hash)
-            if blob:
-                self.put_cache(expect_hash, blob)
-                return blob
+        slot, first = self.reserve_cache(expect_hash)
+        #logging.debug("%s %s %s" % (slot, first, expect_hash))
+
+        if not first:
+            v = slot.get()
+            return v
+
+        try:
+            for service_root in self.shuffled_service_roots(expect_hash):
+                url = service_root + expect_hash
+                api_token = config.get('ARVADOS_API_TOKEN')
+                headers = {'Authorization': "OAuth2 %s" % api_token,
+                           'Accept': 'application/octet-stream'}
+                blob = self.get_url(url, headers, expect_hash)
+                if blob:
+                    slot.set(blob)
+                    self.cap_cache()
+                    return blob
+
+            for location_hint in re.finditer(r'\+K@([a-z0-9]+)', locator):
+                instance = location_hint.group(1)
+                url = 'http://keep.' + instance + '.arvadosapi.com/' + expect_hash
+                blob = self.get_url(url, {}, expect_hash)
+                if blob:
+                    slot.set(blob)
+                    self.cap_cache()
+                    return blob
+        except:
+            slot.set(None)
+            self.cap_cache()
+            raise
+
+        slot.set(None)
+        self.cap_cache()
         raise arvados.errors.NotFoundError("Block not found: %s" % expect_hash)
 
     def get_url(self, url, headers, expect_hash):
         h = httplib2.Http()
         try:
-            resp, content = h.request(url.encode('utf-8'), 'GET',
-                                      headers=headers)
+            logging.info("Request: GET %s" % (url))
+            with timer.Timer() as t:
+                resp, content = h.request(url.encode('utf-8'), 'GET',
+                                          headers=headers)
+            logging.info("Received %s bytes in %s msec (%s MiB/sec)" % (len(content), 
+                                                                        t.msecs, 
+                                                                        (len(content)/(1024*1024))/t.secs))
             if re.match(r'^2\d\d$', resp['status']):
                 m = hashlib.new('md5')
                 m.update(content)
@@ -246,38 +324,6 @@ class KeepClient(object):
             "Write fail for %s: wanted %d but wrote %d" %
             (data_hash, want_copies, have_copies))
 
-    def put_cache(self, locator, data):
-        """Put a block into the cache."""
-        if self.check_cache(locator) != None:
-            return
-        self._cache_lock.acquire()
-        try:
-            # first check cache size and delete stuff from the end if necessary
-            sm = sum([len(a[1]) for a in self._cache]) + len(data)
-            while sm > self._cache_max:
-               print sm, self._cache_max
-                del self._cache[-1]
-                sm = sum([len(a[1]) for a in self._cache]) + len(data)
-
-            # now add the new block at the front of the list
-            self._cache.insert(0, [locator, data])
-        finally:
-            self._cache_lock.release()
-
-    def check_cache(self, locator):
-        """Get a block from the cache.  Also moves the block to the front of the list."""
-        self._cache_lock.acquire()
-        try:
-            for i in xrange(0, len(self._cache)):
-                if self._cache[i][0] == locator:
-                    n = self._cache[i]
-                    del self._cache[i]
-                    self._cache.insert(0, n)
-                    return n[1]   
-        finally:
-            self._cache_lock.release()
-        return None            
-
     @staticmethod
     def sign_for_old_server(data_hash, data):
         return (("-----BEGIN PGP SIGNED MESSAGE-----\n\n\n%d %s\n-----BEGIN PGP SIGNATURE-----\n\n-----END PGP SIGNATURE-----\n" % (int(time.time()), data_hash)) + data)