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