wrw/session.py

   1 import threading, time, pickle, random, os
   2 <<<<<<< HEAD
   3 import cookie, env
   4 =======
   5 from . import cookie, env, proto
   6 >>>>>>> master
   7
   8 __all__ = ["db", "get"]
   9
  10 def gennonce(length):
  11     return os.urandom(length)
  12
  13 class session(object):
  14     def __init__(self, lock, expire=86400 * 7):
  15         self.id = proto.enhex(gennonce(16))
  16         self.dict = {}
  17         self.lock = lock
  18         self.ctime = self.atime = self.mtime = int(time.time())
  19         self.expire = expire
  20         self.dctl = set()
  21         self.dirtyp = False
  22
  23     def dirty(self):
  24         for d in self.dctl:
  25             if d.sessdirty():
  26                 return True
  27         return self.dirtyp
  28
  29     def frozen(self):
  30         for d in self.dctl:
  31             d.sessfrozen()
  32         self.dirtyp = False
  33
  34     def __getitem__(self, key):
  35         return self.dict[key]
  36
  37     def get(self, key, default=None):
  38         return self.dict.get(key, default)
  39
  40     def __setitem__(self, key, value):
  41         self.dict[key] = value
  42         if hasattr(value, "sessdirty"):
  43             self.dctl.add(value)
  44         else:
  45             self.dirtyp = True
  46
  47     def __delitem__(self, key):
  48         old = self.dict.pop(key)
  49         if old in self.dctl:
  50             self.dctl.remove(old)
  51         self.dirtyp = True
  52
  53     def __contains__(self, key):
  54         return key in self.dict
  55
  56     def __getstate__(self):
  57         ret = []
  58         for k, v in self.__dict__.items():
  59             if k == "lock": continue
  60             ret.append((k, v))
  61         return ret
  62
  63     def __setstate__(self, st):
  64         for k, v in st:
  65             self.__dict__[k] = v
  66         # The proper lock is set by the thawer
  67
  68     def __repr__(self):
  69         return "<session %s>" % self.id
  70
  71 class db(object):
  72     def __init__(self, backdb=None, cookiename="wrwsess", path="/"):
  73         self.live = {}
  74         self.cookiename = cookiename
  75         self.path = path
  76         self.lock = threading.Lock()
  77         self.cthread = None
  78         self.freezetime = 3600
  79         self.backdb = backdb
  80
  81     def clean(self):
  82         now = int(time.time())
  83         with self.lock:
  84             clist = self.live.keys()
  85         for sessid in clist:
  86             with self.lock:
  87                 try:
  88                     entry = self.live[sessid]
  89                 except KeyError:
  90                     continue
  91             with entry[0]:
  92                 rm = False
  93                 if entry[1] == "retired":
  94                     pass
  95                 elif entry[1] is None:
  96                     pass
  97                 else:
  98                     sess = entry[1]
  99                     if sess.atime + self.freezetime < now:
 100                         try:
 101                             if sess.dirty():
 102                                 self.freeze(sess)
 103                         except:
 104                             if sess.atime + sess.expire < now:
 105                                 rm = True
 106                         else:
 107                             rm = True
 108                 if rm:
 109                     entry[1] = "retired"
 110                     with self.lock:
 111                         del self.live[sessid]
 112
 113     def cleanloop(self):
 114         try:
 115             while True:
 116                 time.sleep(300)
 117                 self.clean()
 118                 if len(self.live) == 0:
 119                     break
 120         finally:
 121             with self.lock:
 122                 self.cthread = None
 123
 124     def _fetch(self, sessid):
 125         while True:
 126             now = int(time.time())
 127             with self.lock:
 128                 if sessid in self.live:
 129                     entry = self.live[sessid]
 130                 else:
 131                     entry = self.live[sessid] = [threading.RLock(), None]
 132             with entry[0]:
 133                 if isinstance(entry[1], session):
 134                     entry[1].atime = now
 135                     return entry[1]
 136                 elif entry[1] == "retired":
 137                     continue
 138                 elif entry[1] is None:
 139                     try:
 140                         thawed = self.thaw(sessid)
 141                         if thawed.atime + thawed.expire < now:
 142                             raise KeyError()
 143                         thawed.lock = entry[0]
 144                         thawed.atime = now
 145                         entry[1] = thawed
 146                         return thawed
 147                     finally:
 148                         if entry[1] is None:
 149                             entry[1] = "retired"
 150                             with self.lock:
 151                                 del self.live[sessid]
 152                 else:
 153                     raise Exception("Illegal session entry: " + repr(entry[1]))
 154
 155     def checkclean(self):
 156         with self.lock:
 157             if self.cthread is None:
 158                 self.cthread = threading.Thread(target = self.cleanloop)
 159                 self.cthread.setDaemon(True)
 160                 self.cthread.start()
 161
 162     def mksession(self, req):
 163         return session(threading.RLock())
 164
 165     def mkcookie(self, req, sess):
 166         cookie.add(req, self.cookiename, sess.id,
 167                    path=self.path,
 168                    expires=cookie.cdate(time.time() + sess.expire))
 169
 170     def fetch(self, req):
 171         now = int(time.time())
 172         sessid = cookie.get(req, self.cookiename)
 173         new = False
 174         try:
 175             if sessid is None:
 176                 raise KeyError()
 177             sess = self._fetch(sessid)
 178         except KeyError:
 179             sess = self.mksession(req)
 180             new = True
 181
 182         def ckfreeze(req):
 183             if sess.dirty():
 184                 if new:
 185                     self.mkcookie(req, sess)
 186                     with self.lock:
 187                         self.live[sess.id] = [sess.lock, sess]
 188                 try:
 189                     self.freeze(sess)
 190                 except:
 191                     pass
 192                 self.checkclean()
 193         req.oncommit(ckfreeze)
 194         return sess
 195
 196     def thaw(self, sessid):
 197         if self.backdb is None:
 198             raise KeyError()
 199         data = self.backdb[sessid]
 200         try:
 201             return pickle.loads(data)
 202         except Exception, e:
 203             raise KeyError()
 204
 205     def freeze(self, sess):
 206         if self.backdb is None:
 207             raise TypeError()
 208         with sess.lock:
 209             data = pickle.dumps(sess, -1)
 210         self.backdb[sess.id] = data
 211         sess.frozen()
 212
 213     def get(self, req):
 214         return req.item(self.fetch)
 215
 216 class dirback(object):
 217     def __init__(self, path):
 218         self.path = path
 219
 220     def __getitem__(self, key):
 221         try:
 222             with open(os.path.join(self.path, key)) as inf:
 223                 return inf.read()
 224         except IOError:
 225             raise KeyError(key)
 226
 227     def __setitem__(self, key, value):
 228         if not os.path.exists(self.path):
 229             os.makedirs(self.path)
 230         with open(os.path.join(self.path, key), "w") as out:
 231             out.write(value)
 232
 233 default = env.var(db(backdb=dirback(os.path.join("/tmp", "wrwsess-" + str(os.getuid())))))
 234
 235 def get(req):
 236     return default.val.get(req)