22 global_client_object = None
30 def global_client_object():
31 global global_client_object
32 if global_client_object == None:
33 global_client_object = KeepClient()
34 return global_client_object
37 def get(locator, **kwargs):
38 return Keep.global_client_object().get(locator, **kwargs)
41 def put(data, **kwargs):
42 return Keep.global_client_object().put(data, **kwargs)
44 class KeepClient(object):
46 class ThreadLimiter(object):
48 Limit the number of threads running at a given time to
49 {desired successes} minus {successes reported}. When successes
50 reported == desired, wake up the remaining threads and tell
53 Should be used in a "with" block.
55 def __init__(self, todo):
58 self._todo_lock = threading.Semaphore(todo)
59 self._done_lock = threading.Lock()
62 self._todo_lock.acquire()
65 def __exit__(self, type, value, traceback):
66 self._todo_lock.release()
68 def shall_i_proceed(self):
70 Return true if the current thread should do stuff. Return
71 false if the current thread should just stop.
74 return (self._done < self._todo)
76 def increment_done(self):
78 Report that the current thread was successful.
85 Return how many successes were reported.
90 class KeepWriterThread(threading.Thread):
92 Write a blob of data to the given Keep server. Call
93 increment_done() of the given ThreadLimiter if the write
96 def __init__(self, **kwargs):
97 super(KeepClient.KeepWriterThread, self).__init__()
101 with self.args['thread_limiter'] as limiter:
102 if not limiter.shall_i_proceed():
103 # My turn arrived, but the job has been done without
106 logging.debug("KeepWriterThread %s proceeding %s %s" %
107 (str(threading.current_thread()),
108 self.args['data_hash'],
109 self.args['service_root']))
111 url = self.args['service_root'] + self.args['data_hash']
112 api_token = config.get('ARVADOS_API_TOKEN')
113 headers = {'Authorization': "OAuth2 %s" % api_token}
115 resp, content = h.request(url.encode('utf-8'), 'PUT',
117 body=self.args['data'])
118 if (resp['status'] == '401' and
119 re.match(r'Timestamp verification failed', content)):
120 body = KeepClient.sign_for_old_server(
121 self.args['data_hash'],
124 resp, content = h.request(url.encode('utf-8'), 'PUT',
127 if re.match(r'^2\d\d$', resp['status']):
128 logging.debug("KeepWriterThread %s succeeded %s %s" %
129 (str(threading.current_thread()),
130 self.args['data_hash'],
131 self.args['service_root']))
132 return limiter.increment_done()
133 logging.warning("Request fail: PUT %s => %s %s" %
134 (url, resp['status'], content))
135 except (httplib2.HttpLib2Error, httplib.HTTPException) as e:
136 logging.warning("Request fail: PUT %s => %s: %s" %
137 (url, type(e), str(e)))
140 self.lock = threading.Lock()
141 self.service_roots = None
142 self._cache_lock = threading.Lock()
144 # default 256 megabyte cache
145 self.cache_max = 256 * 1024 * 1024
147 def shuffled_service_roots(self, hash):
148 if self.service_roots == None:
151 keep_disks = arvados.api().keep_disks().list().execute()['items']
152 roots = (("http%s://%s:%d/" %
153 ('s' if f['service_ssl_flag'] else '',
157 self.service_roots = sorted(set(roots))
158 logging.debug(str(self.service_roots))
162 # Build an ordering with which to query the Keep servers based on the
163 # contents of the hash.
164 # "hash" is a hex-encoded number at least 8 digits
167 # seed used to calculate the next keep server from 'pool'
168 # to be added to 'pseq'
171 # Keep servers still to be added to the ordering
172 pool = self.service_roots[:]
174 # output probe sequence
177 # iterate while there are servers left to be assigned
180 # ran out of digits in the seed
181 if len(pseq) < len(hash) / 4:
182 # the number of servers added to the probe sequence is less
183 # than the number of 4-digit slices in 'hash' so refill the
184 # seed with the last 4 digits and then append the contents
186 seed = hash[-4:] + hash
188 # refill the seed with the contents of 'hash'
191 # Take the next 8 digits (32 bytes) and interpret as an integer,
192 # then modulus with the size of the remaining pool to get the next
194 probe = int(seed[0:8], 16) % len(pool)
196 print seed[0:8], int(seed[0:8], 16), len(pool), probe
198 # Append the selected server to the probe sequence and remove it
200 pseq += [pool[probe]]
201 pool = pool[:probe] + pool[probe+1:]
203 # Remove the digits just used from the seed
205 logging.debug(str(pseq))
208 class CacheSlot(object):
209 def __init__(self, locator):
210 self.locator = locator
211 self.ready = threading.Event()
218 def set(self, value):
223 if self.content == None:
226 return len(self.content)
229 '''Cap the cache size to self.cache_max'''
230 self._cache_lock.acquire()
232 self._cache = filter(lambda c: not (c.ready.is_set() and c.content == None), self._cache)
233 sm = sum([slot.size() for slot in self._cache])
234 while sm > self.cache_max:
236 sm = sum([slot.size() for a in self._cache])
238 self._cache_lock.release()
240 def reserve_cache(self, locator):
241 '''Reserve a cache slot for the specified locator,
242 or return the existing slot.'''
243 self._cache_lock.acquire()
245 # Test if the locator is already in the cache
246 for i in xrange(0, len(self._cache)):
247 if self._cache[i].locator == locator:
250 # move it to the front
252 self._cache.insert(0, n)
255 # Add a new cache slot for the locator
256 n = KeepClient.CacheSlot(locator)
257 self._cache.insert(0, n)
260 self._cache_lock.release()
262 def get(self, locator):
263 #logging.debug("Keep.get %s" % (locator))
265 if re.search(r',', locator):
266 return ''.join(self.get(x) for x in locator.split(','))
267 if 'KEEP_LOCAL_STORE' in os.environ:
268 return KeepClient.local_store_get(locator)
269 expect_hash = re.sub(r'\+.*', '', locator)
271 slot, first = self.reserve_cache(expect_hash)
272 #logging.debug("%s %s %s" % (slot, first, expect_hash))
279 for service_root in self.shuffled_service_roots(expect_hash):
280 url = service_root + expect_hash
281 api_token = config.get('ARVADOS_API_TOKEN')
282 headers = {'Authorization': "OAuth2 %s" % api_token,
283 'Accept': 'application/octet-stream'}
284 blob = self.get_url(url, headers, expect_hash)
290 for location_hint in re.finditer(r'\+K@([a-z0-9]+)', locator):
291 instance = location_hint.group(1)
292 url = 'http://keep.' + instance + '.arvadosapi.com/' + expect_hash
293 blob = self.get_url(url, {}, expect_hash)
305 raise arvados.errors.NotFoundError("Block not found: %s" % expect_hash)
307 def get_url(self, url, headers, expect_hash):
310 logging.info("Request: GET %s" % (url))
311 with timer.Timer() as t:
312 resp, content = h.request(url.encode('utf-8'), 'GET',
314 logging.info("Received %s bytes in %s msec (%s MiB/sec)" % (len(content),
316 (len(content)/(1024*1024))/t.secs))
317 if re.match(r'^2\d\d$', resp['status']):
318 m = hashlib.new('md5')
321 if md5 == expect_hash:
323 logging.warning("Checksum fail: md5(%s) = %s" % (url, md5))
324 except Exception as e:
325 logging.info("Request fail: GET %s => %s: %s" %
326 (url, type(e), str(e)))
329 def put(self, data, **kwargs):
330 if 'KEEP_LOCAL_STORE' in os.environ:
331 return KeepClient.local_store_put(data)
332 m = hashlib.new('md5')
334 data_hash = m.hexdigest()
336 want_copies = kwargs.get('copies', 2)
337 if not (want_copies > 0):
340 thread_limiter = KeepClient.ThreadLimiter(want_copies)
341 for service_root in self.shuffled_service_roots(data_hash):
342 t = KeepClient.KeepWriterThread(data=data,
344 service_root=service_root,
345 thread_limiter=thread_limiter)
350 have_copies = thread_limiter.done()
351 if have_copies == want_copies:
352 return (data_hash + '+' + str(len(data)))
353 raise arvados.errors.KeepWriteError(
354 "Write fail for %s: wanted %d but wrote %d" %
355 (data_hash, want_copies, have_copies))
358 def sign_for_old_server(data_hash, data):
359 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)
363 def local_store_put(data):
364 m = hashlib.new('md5')
367 locator = '%s+%d' % (md5, len(data))
368 with open(os.path.join(os.environ['KEEP_LOCAL_STORE'], md5 + '.tmp'), 'w') as f:
370 os.rename(os.path.join(os.environ['KEEP_LOCAL_STORE'], md5 + '.tmp'),
371 os.path.join(os.environ['KEEP_LOCAL_STORE'], md5))
375 def local_store_get(locator):
376 r = re.search('^([0-9a-f]{32,})', locator)
378 raise arvados.errors.NotFoundError(
379 "Invalid data locator: '%s'" % locator)
380 if r.group(0) == config.EMPTY_BLOCK_LOCATOR.split('+')[0]:
382 with open(os.path.join(os.environ['KEEP_LOCAL_STORE'], r.group(0)), 'r') as f: