Add support for ARVADOS_API_PORT.
[arvados.git] / sdk / python / arvados.py
index 166fbd14ef9632030eb5dacd57c2d2159f88ce4e..dfe9ff65306e5d9a506011ed3a879d910d576348 100644 (file)
@@ -16,6 +16,7 @@ import bz2
 import zlib
 import fcntl
 import time
+import threading
 
 from apiclient import errors
 from apiclient.discovery import build
@@ -23,7 +24,7 @@ from apiclient.discovery import build
 if 'ARVADOS_DEBUG' in os.environ:
     logging.basicConfig(level=logging.DEBUG)
 
-class CredentialsFromEnv:
+class CredentialsFromEnv(object):
     @staticmethod
     def http_request(self, uri, **kwargs):
         from httplib import BadStatusLine
@@ -45,8 +46,10 @@ class CredentialsFromEnv:
         http.request = types.MethodType(self.http_request, http)
         return http
 
-url = ('https://%s/discovery/v1/apis/'
-       '{api}/{apiVersion}/rest' % os.environ['ARVADOS_API_HOST'])
+url = ('https://%s:%s/discovery/v1/apis/'
+       '{api}/{apiVersion}/rest' %
+           (os.environ['ARVADOS_API_HOST'],
+            os.environ.get('ARVADOS_API_PORT') or "443"))
 credentials = CredentialsFromEnv()
 
 # Use system's CA certificates (if we find them) instead of httplib2's
@@ -95,7 +98,7 @@ def current_job():
 def api():
     return service
 
-class JobTask:
+class JobTask(object):
     def __init__(self, parameters=dict(), runtime_constraints=dict()):
         print "init jobtask %s %s" % (parameters, runtime_constraints)
 
@@ -431,30 +434,7 @@ class util:
                 allfiles += [ent_base]
         return allfiles
 
-class DataReader:
-    def __init__(self, data_locator):
-        self.data_locator = data_locator
-        self.p = subprocess.Popen(["whget", "-r", self.data_locator, "-"],
-                                  stdout=subprocess.PIPE,
-                                  stdin=None, stderr=subprocess.PIPE,
-                                  shell=False, close_fds=True)
-    def __enter__(self):
-        pass
-    def __exit__(self):
-        self.close()
-    def read(self, size, **kwargs):
-        return self.p.stdout.read(size, **kwargs)
-    def close(self):
-        self.p.stdout.close()
-        if not self.p.stderr.closed:
-            for err in self.p.stderr:
-                print >> sys.stderr, err
-            self.p.stderr.close()
-        self.p.wait()
-        if self.p.returncode != 0:
-            raise Exception("whget subprocess exited %d" % self.p.returncode)
-
-class StreamFileReader:
+class StreamFileReader(object):
     def __init__(self, stream, pos, size, name):
         self._stream = stream
         self._pos = pos
@@ -526,7 +506,7 @@ class StreamFileReader:
         return string.join(self._stream.tokens_for_range(self._pos, self._size),
                            " ") + "\n"
 
-class StreamReader:
+class StreamReader(object):
     def __init__(self, tokens):
         self._tokens = tokens
         self._current_datablock_data = None
@@ -638,7 +618,7 @@ class StreamReader:
         self._pos += len(data)
         return data
 
-class CollectionReader:
+class CollectionReader(object):
     def __init__(self, manifest_locator_or_text):
         if re.search(r'^\S+( [a-f0-9]{32,}(\+\S+)*)+( \d+:\d+:\S+)+\n', manifest_locator_or_text):
             self._manifest_text = manifest_locator_or_text
@@ -675,7 +655,7 @@ class CollectionReader:
         self._populate()
         return self._manifest_text
 
-class CollectionWriter:
+class CollectionWriter(object):
     KEEP_BLOCK_SIZE = 2**26
     def __init__(self):
         self._data_buffer = []
@@ -803,19 +783,109 @@ class Keep:
         return global_client_object
 
     @staticmethod
-    def get(locator):
-        return Keep.global_client_object().get(locator)
+    def get(locator, **kwargs):
+        return Keep.global_client_object().get(locator, **kwargs)
 
     @staticmethod
-    def put(data):
-        return Keep.global_client_object().put(data)
+    def put(data, **kwargs):
+        return Keep.global_client_object().put(data, **kwargs)
+
+class KeepClient(object):
+
+    class ThreadLimiter(object):
+        """
+        Limit the number of threads running at a given time to
+        {desired successes} minus {successes reported}. When successes
+        reported == desired, wake up the remaining threads and tell
+        them to quit.
+
+        Should be used in a "with" block.
+        """
+        def __init__(self, todo):
+            self._todo = todo
+            self._done = 0
+            self._todo_lock = threading.Semaphore(todo)
+            self._done_lock = threading.Lock()
+        def __enter__(self):
+            self._todo_lock.acquire()
+            return self
+        def __exit__(self, type, value, traceback):
+            self._todo_lock.release()
+        def shall_i_proceed(self):
+            """
+            Return true if the current thread should do stuff. Return
+            false if the current thread should just stop.
+            """
+            with self._done_lock:
+                return (self._done < self._todo)
+        def increment_done(self):
+            """
+            Report that the current thread was successful.
+            """
+            with self._done_lock:
+                self._done += 1
+        def done(self):
+            """
+            Return how many successes were reported.
+            """
+            with self._done_lock:
+                return self._done
+
+    class KeepWriterThread(threading.Thread):
+        """
+        Write a blob of data to the given Keep server. Call
+        increment_done() of the given ThreadLimiter if the write
+        succeeds.
+        """
+        def __init__(self, **kwargs):
+            super(KeepClient.KeepWriterThread, self).__init__()
+            self.args = kwargs
+        def run(self):
+            with self.args['thread_limiter'] as limiter:
+                if not limiter.shall_i_proceed():
+                    # My turn arrived, but the job has been done without
+                    # me.
+                    return
+                logging.debug("KeepWriterThread %s proceeding %s %s" %
+                              (str(threading.current_thread()),
+                               self.args['data_hash'],
+                               self.args['service_root']))
+                h = httplib2.Http()
+                url = self.args['service_root'] + self.args['data_hash']
+                api_token = os.environ['ARVADOS_API_TOKEN']
+                headers = {'Authorization': "OAuth2 %s" % api_token}
+                try:
+                    resp, content = h.request(url.encode('utf-8'), 'PUT',
+                                              headers=headers,
+                                              body=self.args['data'])
+                    if (resp['status'] == '401' and
+                        re.match(r'Timestamp verification failed', content)):
+                        body = KeepClient.sign_for_old_server(
+                            self.args['data_hash'],
+                            self.args['data'])
+                        h = httplib2.Http()
+                        resp, content = h.request(url.encode('utf-8'), 'PUT',
+                                                  headers=headers,
+                                                  body=body)
+                    if re.match(r'^2\d\d$', resp['status']):
+                        logging.debug("KeepWriterThread %s succeeded %s %s" %
+                                      (str(threading.current_thread()),
+                                       self.args['data_hash'],
+                                       self.args['service_root']))
+                        return limiter.increment_done()
+                    logging.warning("Request fail: PUT %s => %s %s" %
+                                    (url, resp['status'], content))
+                except (httplib2.HttpLib2Error, httplib.HTTPException) as e:
+                    logging.warning("Request fail: PUT %s => %s: %s" %
+                                    (url, type(e), str(e)))
 
-class KeepClient:
     def __init__(self):
+        self.lock = threading.Lock()
         self.service_roots = None
 
     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 '',
@@ -824,6 +894,7 @@ class KeepClient:
                      for f in keep_disks)
             self.service_roots = sorted(set(roots))
             logging.debug(str(self.service_roots))
+            self.lock.release()
         seed = hash
         pool = self.service_roots[:]
         pseq = []
@@ -851,7 +922,8 @@ class KeepClient:
             headers = {'Authorization': "OAuth2 %s" % api_token,
                        'Accept': 'application/octet-stream'}
             try:
-                resp, content = h.request(url, 'GET', headers=headers)
+                resp, content = h.request(url.encode('utf-8'), 'GET',
+                                          headers=headers)
                 if re.match(r'^2\d\d$', resp['status']):
                     m = hashlib.new('md5')
                     m.update(content)
@@ -872,36 +944,27 @@ class KeepClient:
         data_hash = m.hexdigest()
         have_copies = 0
         want_copies = kwargs.get('copies', 2)
+        if not (want_copies > 0):
+            return data_hash
+        threads = []
+        thread_limiter = KeepClient.ThreadLimiter(want_copies)
         for service_root in self.shuffled_service_roots(data_hash):
-            h = httplib2.Http()
-            url = service_root + data_hash
-            api_token = os.environ['ARVADOS_API_TOKEN']
-            headers = {'Authorization': "OAuth2 %s" % api_token}
-            try:
-                resp, content = h.request(url, 'PUT',
-                                          headers=headers,
-                                          body=data)
-                if (resp['status'] == '401' and
-                    re.match(r'Timestamp verification failed', content)):
-                    body = self.sign_for_old_server(data_hash, data)
-                    h = httplib2.Http()
-                    resp, content = h.request(url, 'PUT',
-                                              headers=headers,
-                                              body=body)
-                if re.match(r'^2\d\d$', resp['status']):
-                    have_copies += 1
-                    if have_copies == want_copies:
-                        return data_hash + '+' + str(len(data))
-                else:
-                    logging.warning("Request fail: PUT %s => %s %s" %
-                                    (url, resp['status'], content))
-            except (httplib2.HttpLib2Error, httplib.HTTPException) as e:
-                logging.warning("Request fail: PUT %s => %s: %s" %
-                                (url, type(e), str(e)))
+            t = KeepClient.KeepWriterThread(data=data,
+                                            data_hash=data_hash,
+                                            service_root=service_root,
+                                            thread_limiter=thread_limiter)
+            t.start()
+            threads += [t]
+        for t in threads:
+            t.join()
+        have_copies = thread_limiter.done()
+        if have_copies == want_copies:
+            return (data_hash + '+' + str(len(data)))
         raise Exception("Write fail for %s: wanted %d but wrote %d" %
                         (data_hash, want_copies, have_copies))
 
-    def sign_for_old_server(self, data_hash, data):
+    @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)