didex/cache.py

   1 import threading, weakref
   2
   3 class entry(object):
   4     __slots__ = ["p", "n", "id", "obj", "st", "lk"]
   5     def __init__(self, id, c):
   6         self.id = id
   7         self.obj = None
   8         self.st = None
   9         self.lk = None
  10         self.n = c.mru
  11         self.p = None
  12         if c.mru is not None:
  13             c.mru.p = self
  14             c.mru = self
  15         else:
  16             c.mru = c.lru = self
  17         c.n += 1
  18
  19     def relink(self, c):
  20         if c.mru is self:
  21             return
  22         if self.n is not None:
  23             self.n.p = self.p
  24         self.p.n = self.n
  25         if c.lru is self:
  26             c.lru = self.p
  27         self.p = None
  28         self.n = c.mru
  29         c.mru.p = self
  30         c.mru = self
  31
  32     def remove(self, c):
  33         if self.n is not None:
  34             self.n.p = self.p
  35         if self.p is not None:
  36             self.p.n = self.n
  37         if c.mru is self:
  38             c.mru = self.n
  39         if c.lru is self:
  40             c.lru = self.p
  41         c.n -= 1
  42
  43 class cache(object):
  44     def __init__(self, *, keep=1000):
  45         self.keep = keep
  46         self.cur = {}
  47         self.mru = self.lru = None
  48         self.n = 0
  49         self.lk = threading.Lock()
  50
  51     def _trim(self, n):
  52         ent = self.lru
  53         for i in range(self.n - n):
  54             if ent.st == "l":
  55                 ent.obj = weakref.ref(ent.obj)
  56                 ent.st = "w"
  57             elif ent.st == "w" and ent.obj() is None:
  58                 del self.cur[ent.id]
  59                 ent.remove(self)
  60                 ent.st = "r"
  61             ent = ent.p
  62
  63     def get(self, id, load=True):
  64         while True:
  65             with self.lk:
  66                 ent = self.cur.get(id)
  67                 if ent is None:
  68                     if not load:
  69                         raise KeyError(id)
  70                     self.cur[id] = ent = entry(id, self)
  71                     ent.lk = lk = threading.Lock()
  72                     ent.st = "ld"
  73                     st = None
  74                     self._trim(self.keep)
  75                 elif ent.st == "l":
  76                     ent.relink(self)
  77                     return ent.obj
  78                 elif ent.st == "w":
  79                     ret = ent.obj()
  80                     if ret is None:
  81                         del self.cur[id]
  82                         ent.remove(self)
  83                         ent.st = "r"
  84                         continue
  85                     return ret
  86                 elif ent.st == "ld":
  87                     lk = ent.lk
  88                     st = "ld"
  89                     if lk is None:
  90                         continue
  91                 elif ent.st == "r":
  92                     continue
  93             with lk:
  94                 if st is None:
  95                     try:
  96                         ret = ent.obj = self.load(id)
  97                         ent.st = "l"
  98                         return ret
  99                     except:
 100                         with self.lk:
 101                             del self.cur[id]
 102                             ent.remove(self)
 103                             ent.st = "r"
 104                         raise
 105                     finally:
 106                         ent.lk = None
 107                 elif st == "ld":
 108                     continue
 109
 110     def put(self, id, ob):
 111         while True:
 112             with self.lk:
 113                 ent = self.cur.get(id)
 114                 if ent is None:
 115                     self.cur[id] = ent = entry(id, self)
 116                     ent.obj = ob
 117                     ent.st = "l"
 118                     self._trim(self.keep)
 119                     return
 120                 elif ent.st == "l":
 121                     ent.obj = ob
 122                     return
 123                 elif ent.st == "w":
 124                     ent.obj = ob
 125                     return
 126                 elif ent.st == "r":
 127                     continue
 128                 elif ent.st == "ld":
 129                     lk = ent.lk
 130                     if lk is None:
 131                         continue
 132             with lk:
 133                 continue
 134
 135     def remove(self, id):
 136         while True:
 137             with self.lk:
 138                 ent = self.cur.get(id)
 139                 if ent is None:
 140                     return
 141                 elif ent.st == "ld":
 142                     lk = ent.lk
 143                     if lk is None:
 144                         continue
 145                 else:
 146                     del self.cur[id]
 147                     ent.remove(self)
 148                     ent.st = "r"
 149                     return
 150             with lk:
 151                 continue
 152
 153     def load(self, id):
 154         raise KeyError()