forked from grantjenks/python-diskcache
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathrecipes.py
More file actions
570 lines (452 loc) · 16.8 KB
/
recipes.py
File metadata and controls
570 lines (452 loc) · 16.8 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
"""Disk Cache Recipes
"""
import functools
import math
import os
import random
import threading
import time
from .core import ENOVAL, args_to_key, full_name
class Averager:
"""Recipe for calculating a running average.
Sometimes known as "online statistics," the running average maintains the
total and count. The average can then be calculated at any time.
>>> import diskcache
>>> cache = diskcache.FanoutCache()
>>> ave = Averager(cache, 'latency')
>>> ave.add(0.080)
>>> ave.add(0.120)
>>> ave.get()
0.1
>>> ave.add(0.160)
>>> ave.pop()
0.12
>>> print(ave.get())
None
"""
def __init__(self, cache, key, expire=None, tag=None):
self._cache = cache
self._key = key
self._expire = expire
self._tag = tag
def add(self, value):
"Add `value` to average."
with self._cache.transact(retry=True):
total, count = self._cache.get(self._key, default=(0.0, 0))
total += value
count += 1
self._cache.set(
self._key,
(total, count),
expire=self._expire,
tag=self._tag,
)
def get(self):
"Get current average or return `None` if count equals zero."
total, count = self._cache.get(self._key, default=(0.0, 0), retry=True)
return None if count == 0 else total / count
def pop(self):
"Return current average and delete key."
total, count = self._cache.pop(self._key, default=(0.0, 0), retry=True)
return None if count == 0 else total / count
class Lock:
"""Recipe for cross-process and cross-thread lock.
>>> import diskcache
>>> cache = diskcache.Cache()
>>> lock = Lock(cache, 'report-123')
>>> lock.acquire()
>>> lock.release()
>>> with lock:
... pass
"""
def __init__(self, cache, key, expire=None, tag=None):
self._cache = cache
self._key = key
self._expire = expire
self._tag = tag
def acquire(self):
"Acquire lock using spin-lock algorithm."
while True:
added = self._cache.add(
self._key,
None,
expire=self._expire,
tag=self._tag,
retry=True,
)
if added:
break
time.sleep(0.001)
def release(self):
"Release lock by deleting key."
self._cache.delete(self._key, retry=True)
def locked(self):
"Return true if the lock is acquired."
return self._key in self._cache
def __enter__(self):
self.acquire()
def __exit__(self, *exc_info):
self.release()
class RLock:
"""Recipe for cross-process and cross-thread re-entrant lock.
>>> import diskcache
>>> cache = diskcache.Cache()
>>> rlock = RLock(cache, 'user-123')
>>> rlock.acquire()
>>> rlock.acquire()
>>> rlock.release()
>>> with rlock:
... pass
>>> rlock.release()
>>> rlock.release()
Traceback (most recent call last):
...
AssertionError: cannot release un-acquired lock
"""
def __init__(self, cache, key, expire=None, tag=None):
self._cache = cache
self._key = key
self._expire = expire
self._tag = tag
def acquire(self):
"Acquire lock by incrementing count using spin-lock algorithm."
pid = os.getpid()
tid = threading.get_ident()
pid_tid = '{}-{}'.format(pid, tid)
while True:
with self._cache.transact(retry=True):
value, count = self._cache.get(self._key, default=(None, 0))
if pid_tid == value or count == 0:
self._cache.set(
self._key,
(pid_tid, count + 1),
expire=self._expire,
tag=self._tag,
)
return
time.sleep(0.001)
def release(self):
"Release lock by decrementing count."
pid = os.getpid()
tid = threading.get_ident()
pid_tid = '{}-{}'.format(pid, tid)
with self._cache.transact(retry=True):
value, count = self._cache.get(self._key, default=(None, 0))
is_owned = pid_tid == value and count > 0
assert is_owned, 'cannot release un-acquired lock'
self._cache.set(
self._key,
(value, count - 1),
expire=self._expire,
tag=self._tag,
)
def __enter__(self):
self.acquire()
def __exit__(self, *exc_info):
self.release()
class BoundedSemaphore:
"""Recipe for cross-process and cross-thread bounded semaphore.
>>> import diskcache
>>> cache = diskcache.Cache()
>>> semaphore = BoundedSemaphore(cache, 'max-cons', value=2)
>>> semaphore.acquire()
>>> semaphore.acquire()
>>> semaphore.release()
>>> with semaphore:
... pass
>>> semaphore.release()
>>> semaphore.release()
Traceback (most recent call last):
...
AssertionError: cannot release un-acquired semaphore
"""
def __init__(self, cache, key, value=1, expire=None, tag=None):
self._cache = cache
self._key = key
self._value = value
self._expire = expire
self._tag = tag
def acquire(self):
"Acquire semaphore by decrementing value using spin-lock algorithm."
while True:
with self._cache.transact(retry=True):
value = self._cache.get(self._key, default=self._value)
if value > 0:
self._cache.set(
self._key,
value - 1,
expire=self._expire,
tag=self._tag,
)
return
time.sleep(0.001)
def release(self):
"Release semaphore by incrementing value."
with self._cache.transact(retry=True):
value = self._cache.get(self._key, default=self._value)
assert self._value > value, 'cannot release un-acquired semaphore'
value += 1
self._cache.set(
self._key,
value,
expire=self._expire,
tag=self._tag,
)
def __enter__(self):
self.acquire()
def __exit__(self, *exc_info):
self.release()
def throttle(
cache,
count,
seconds,
name=None,
expire=None,
tag=None,
time_func=time.time,
sleep_func=time.sleep,
):
"""Decorator to throttle calls to function.
>>> import diskcache, time
>>> cache = diskcache.Cache()
>>> count = 0
>>> @throttle(cache, 2, 1) # 2 calls per 1 second
... def increment():
... global count
... count += 1
>>> start = time.time()
>>> while (time.time() - start) <= 2:
... increment()
>>> count in (6, 7) # 6 or 7 calls depending on CPU load
True
"""
def decorator(func):
rate = count / float(seconds)
key = full_name(func) if name is None else name
now = time_func()
cache.set(key, (now, count), expire=expire, tag=tag, retry=True)
@functools.wraps(func)
def wrapper(*args, **kwargs):
while True:
with cache.transact(retry=True):
last, tally = cache.get(key)
now = time_func()
tally += (now - last) * rate
delay = 0
if tally > count:
cache.set(key, (now, count - 1), expire)
elif tally >= 1:
cache.set(key, (now, tally - 1), expire)
else:
delay = (1 - tally) / rate
if delay:
sleep_func(delay)
else:
break
return func(*args, **kwargs)
return wrapper
return decorator
def barrier(cache, lock_factory, name=None, expire=None, tag=None):
"""Barrier to calling decorated function.
Supports different kinds of locks: Lock, RLock, BoundedSemaphore.
>>> import diskcache, time
>>> cache = diskcache.Cache()
>>> @barrier(cache, Lock)
... def work(num):
... print('worker started')
... time.sleep(1)
... print('worker finished')
>>> import multiprocessing.pool
>>> pool = multiprocessing.pool.ThreadPool(2)
>>> _ = pool.map(work, range(2))
worker started
worker finished
worker started
worker finished
>>> pool.terminate()
"""
def decorator(func):
key = full_name(func) if name is None else name
lock = lock_factory(cache, key, expire=expire, tag=tag)
@functools.wraps(func)
def wrapper(*args, **kwargs):
with lock:
return func(*args, **kwargs)
return wrapper
return decorator
def memoize_stampede(cache, expire, name=None, typed=False, tag=None, beta=1):
"""Memoizing cache decorator with cache stampede protection.
Cache stampedes are a type of system overload that can occur when parallel
computing systems using memoization come under heavy load. This behaviour
is sometimes also called dog-piling, cache miss storm, cache choking, or
the thundering herd problem.
The memoization decorator implements cache stampede protection through
early recomputation. Early recomputation of function results will occur
probabilistically before expiration in a background thread of
execution. Early probabilistic recomputation is based on research by
Vattani, A.; Chierichetti, F.; Lowenstein, K. (2015), Optimal Probabilistic
Cache Stampede Prevention, VLDB, pp. 886-897, ISSN 2150-8097
If name is set to None (default), the callable name will be determined
automatically.
If typed is set to True, function arguments of different types will be
cached separately. For example, f(3) and f(3.0) will be treated as distinct
calls with distinct results.
The original underlying function is accessible through the `__wrapped__`
attribute. This is useful for introspection, for bypassing the cache, or
for rewrapping the function with a different cache.
>>> from diskcache import Cache
>>> cache = Cache()
>>> @memoize_stampede(cache, expire=1)
... def fib(number):
... if number == 0:
... return 0
... elif number == 1:
... return 1
... else:
... return fib(number - 1) + fib(number - 2)
>>> print(fib(100))
354224848179261915075
An additional `__cache_key__` attribute can be used to generate the cache
key used for the given arguments.
>>> key = fib.__cache_key__(100)
>>> del cache[key]
Remember to call memoize when decorating a callable. If you forget, then a
TypeError will occur.
:param cache: cache to store callable arguments and return values
:param float expire: seconds until arguments expire
:param str name: name given for callable (default None, automatic)
:param bool typed: cache different types separately (default False)
:param str tag: text to associate with arguments (default None)
:return: callable decorator
"""
# Caution: Nearly identical code exists in Cache.memoize
def decorator(func):
"Decorator created by memoize call for callable."
base = (full_name(func),) if name is None else (name,)
def timer(*args, **kwargs):
"Time execution of `func` and return result and time delta."
start = time.time()
result = func(*args, **kwargs)
delta = time.time() - start
return result, delta
@functools.wraps(func)
def wrapper(*args, **kwargs):
"Wrapper for callable to cache arguments and return values."
key = wrapper.__cache_key__(*args, **kwargs)
pair, expire_time = cache.get(
key,
default=ENOVAL,
expire_time=True,
retry=True,
)
if pair is not ENOVAL:
result, delta = pair
now = time.time()
ttl = expire_time - now
if (-delta * beta * math.log(random.random())) < ttl:
return result # Cache hit.
# Check whether a thread has started for early recomputation.
thread_key = key + (ENOVAL,)
thread_added = cache.add(
thread_key,
None,
expire=delta,
retry=True,
)
if thread_added:
# Start thread for early recomputation.
def recompute():
with cache:
pair = timer(*args, **kwargs)
cache.set(
key,
pair,
expire=expire,
tag=tag,
retry=True,
)
thread = threading.Thread(target=recompute)
thread.daemon = True
thread.start()
return result
pair = timer(*args, **kwargs)
cache.set(key, pair, expire=expire, tag=tag, retry=True)
return pair[0]
def __cache_key__(*args, **kwargs):
"Make key for cache given function arguments."
return args_to_key(base, args, kwargs, typed)
wrapper.__cache_key__ = __cache_key__
return wrapper
return decorator
################################################################################
import asyncio
import functools
import types
from diskcache import Cache
class AsyncMeta(type):
def __new__(cls, name, bases, attrs):
base = attrs['__wrapped__']
def __init__(self, loop=None, executor=None):
if loop is None:
loop = asyncio.get_event_loop()
run = functools.partial(loop.run_in_executor, executor)
self.__dict__['_async_meta_run'] = run
self.__dict__['_async_meta_thing'] = None
assert '__init__' not in attrs
attrs['__init__'] = __init__
async def initialize(self, *args, **kwargs):
"""Initialize the thing attribute."""
func = functools.partial(base, *args, **kwargs)
thing = await self._async_meta_run(func)
self.__dict__['_async_meta_thing'] = thing
assert 'initialize' not in attrs
attrs['initialize'] = initialize
def __getattr__(self, name):
return getattr(self._async_meta_thing, name)
assert '__getattr__' not in attrs
attrs['__getattr__'] = __getattr__
def __setattr__(self, name, value):
setattr(self._async_meta_thing, name, value)
assert '__setattr__' not in attrs
attrs['__setattr__'] = __setattr__
def __delattr__(self, name):
delattr(self._async_meta_thing, name)
assert '__delattr__' not in attrs
attrs['__delattr__'] = __delattr__
# TODO: Add support for __aiter__, __anext__, __aenter__, and __aexit__
def make_method(func):
"""Make an async wrapper method."""
@functools.wraps(func)
async def method(self, *args, **kwargs):
"""Async wrapper method."""
# `AsyncThing` wraps `Thing` so pass `self._async_meta_thing` as the
# first argument to be bound to `self` in `Thing` method calls.
call = functools.partial(func, self._async_meta_thing, *args, **kwargs)
return await self._async_meta_run(call)
return method
# Iterate the attributes of the cache and make methods for `AsyncCache`.
for name in dir(base):
if name.startswith('_'):
# Only support the "public" methods.
continue
attr = getattr(base, name)
if not isinstance(attr, types.FunctionType):
continue
# TODO: Add support for contextlib.contextmanager types.
method = make_method(attr)
attrs[name] = method
bases = (object,)
return super().__new__(cls, name, bases, attrs)
class AsyncCache(metaclass=AsyncMeta):
__wrapped__ = Cache
###############################################################################
async def main():
cache = AsyncCache()
await cache.initialize(directory='/tmp/diskcache/async')
assert cache.directory == '/tmp/diskcache/async'
await cache.set('key', 'value')
print(await cache.get('key'))
async with cache.transact():
pass
if __name__ == '__main__':
asyncio.run(main())