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