Inserted __future__ import of with_statement to be Jython (that is, 2.5) compatible.
[wrw.git] / 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)