wrw/session.py

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