callcgi: Fixed possible deadlock problem on aborted requests.
[ashd.git] / python3 / ashd / async.py
1 import sys, os, errno, threading, select, traceback
2
3 class epoller(object):
4     exc_handler = None
5
6     def __init__(self, check=None):
7         self.registered = {}
8         self.fdcache = {}
9         self.lock = threading.RLock()
10         self.ep = None
11         self.th = None
12         self.stopped = False
13         self.loopcheck = set()
14         if check is not None:
15             self.loopcheck.add(check)
16         self._daemon = True
17
18     @staticmethod
19     def _evsfor(ch):
20         return ((select.EPOLLIN if ch.readable else 0) |
21                 (select.EPOLLOUT if ch.writable else 0))
22
23     def _ckrun(self):
24         if self.registered and self.th is None:
25             th = threading.Thread(target=self._run, name="Async epoll thread")
26             th.daemon = self._daemon
27             th.start()
28             self.th = th
29
30     def exception(self, ch, *exc):
31         self.remove(ch)
32         if self.exc_handler is None:
33             traceback.print_exception(*exc)
34         else:
35             self.exc_handler(ch, *exc)
36
37     def _cb(self, ch, nm):
38         try:
39             m = getattr(ch, nm, None)
40             if m is None:
41                 raise AttributeError("%r has no %s method" % (ch, nm))
42             m()
43         except Exception as exc:
44             self.exception(ch, *sys.exc_info())
45
46     def _closeall(self):
47         with self.lock:
48             while self.registered:
49                 fd, (ch, evs) = next(iter(self.registered.items()))
50                 del self.registered[fd]
51                 self.ep.unregister(fd)
52                 self._cb(ch, "close")
53
54     def _run(self):
55         ep = select.epoll()
56         try:
57             with self.lock:
58                 try:
59                     for fd, (ob, evs) in self.registered.items():
60                         ep.register(fd, evs)
61                 except:
62                     self.registered.clear()
63                     raise
64                 self.ep = ep
65
66             while self.registered:
67                 for ck in self.loopcheck:
68                     ck(self)
69                 if self.stopped:
70                     self._closeall()
71                     break
72                 try:
73                     evlist = ep.poll(10)
74                 except IOError as exc:
75                     if exc.errno == errno.EINTR:
76                         continue
77                     raise
78                 for fd, evs in evlist:
79                     with self.lock:
80                         if fd not in self.registered:
81                             continue
82                         ch, cevs = self.registered[fd]
83                         if fd in self.registered and evs & (select.EPOLLIN | select.EPOLLHUP | select.EPOLLERR):
84                             self._cb(ch, "read")
85                         if fd in self.registered and evs & select.EPOLLOUT:
86                             self._cb(ch, "write")
87                         if fd in self.registered:
88                             nevs = self._evsfor(ch)
89                             if nevs == 0:
90                                 del self.fdcache[ch]
91                                 del self.registered[fd]
92                                 ep.unregister(fd)
93                                 self._cb(ch, "close")
94                             elif nevs != cevs:
95                                 self.registered[fd] = ch, nevs
96                                 ep.modify(fd, nevs)
97
98         finally:
99             with self.lock:
100                 self.th = None
101                 self.ep = None
102                 self._ckrun()
103             ep.close()
104
105     @property
106     def daemon(self): return self._daemon
107     @daemon.setter
108     def daemon(self, value):
109         self._daemon = bool(value)
110         with self.lock:
111             if self.th is not None:
112                 self.th = daemon = self._daemon
113
114     def add(self, ch):
115         with self.lock:
116             fd = ch.fileno()
117             if fd in self.registered:
118                 raise KeyError("fd %i is already registered" % fd)
119             evs = self._evsfor(ch)
120             if evs == 0:
121                 ch.close()
122                 return
123             ch.watcher = self
124             self.fdcache[ch] = fd
125             self.registered[fd] = (ch, evs)
126             if self.ep:
127                 self.ep.register(fd, evs)
128             self._ckrun()
129
130     def remove(self, ch, ignore=False):
131         with self.lock:
132             try:
133                 fd = self.fdcache[ch]
134             except KeyError:
135                 if ignore:
136                     return
137                 raise KeyError("fd %i is not registered" % fd)
138             pch, cevs = self.registered[fd]
139             if pch is not ch:
140                 raise ValueError("fd %i registered via object %r, cannot remove with %r" % (pch, ch))
141             del self.fdcache[ch]
142             del self.registered[fd]
143             if self.ep:
144                 self.ep.unregister(fd)
145             ch.close()
146
147     def update(self, ch, ignore=False):
148         with self.lock:
149             try:
150                 fd = self.fdcache[ch]
151             except KeyError:
152                 if ignore:
153                     return
154                 raise KeyError("fd %i is not registered" % fd)
155             pch, cevs = self.registered[fd]
156             if pch is not ch:
157                 raise ValueError("fd %i registered via object %r, cannot update with %r" % (pch, ch))
158             evs = self._evsfor(ch)
159             if evs == 0:
160                 del self.fdcache[ch]
161                 del self.registered[fd]
162                 if self.ep:
163                     self.ep.unregister(fd)
164                 ch.close()
165             elif evs != cevs:
166                 self.registered[fd] = ch, evs
167                 if self.ep:
168                     self.ep.modify(fd, evs)
169
170     def stop(self):
171         if threading.current_thread() == self.th:
172             self.stopped = True
173         else:
174             def tgt():
175                 self.stopped = True
176             cb = callbuffer()
177             cb.call(tgt)
178             cb.stop()
179             self.add(cb)
180
181 def watcher():
182     return epoller()
183
184 class channel(object):
185     readable = False
186     writable = False
187
188     def __init__(self):
189         self.watcher = None
190
191     def fileno(self):
192         raise NotImplementedError("fileno()")
193
194     def close(self):
195         pass
196
197 class sockbuffer(channel):
198     def __init__(self, socket, **kwargs):
199         super().__init__(**kwargs)
200         self.sk = socket
201         self.eof = False
202         self.obuf = bytearray()
203
204     def fileno(self):
205         return self.sk.fileno()
206
207     def close(self):
208         self.sk.close()
209
210     def gotdata(self, data):
211         if data == b"":
212             self.eof = True
213
214     def send(self, data, eof=False):
215         self.obuf.extend(data)
216         if eof:
217             self.eof = True
218         if self.watcher is not None:
219             self.watcher.update(self, True)
220
221     @property
222     def readable(self):
223         return not self.eof
224     def read(self):
225         try:
226             data = self.sk.recv(1024)
227             self.gotdata(data)
228         except IOError:
229             self.obuf[:] = b""
230             self.eof = True
231
232     @property
233     def writable(self):
234         return bool(self.obuf);
235     def write(self):
236         try:
237             ret = self.sk.send(self.obuf)
238             self.obuf[:ret] = b""
239         except IOError:
240             self.obuf[:] = b""
241             self.eof = True
242
243 class callbuffer(channel):
244     def __init__(self, **kwargs):
245         super().__init__(**kwargs)
246         self.queue = []
247         self.rp, self.wp = os.pipe()
248         self.lock = threading.Lock()
249         self.eof = False
250
251     def fileno(self):
252         return self.rp
253
254     def close(self):
255         with self.lock:
256             try:
257                 if self.wp >= 0:
258                     os.close(self.wp)
259                 self.wp = -1
260             finally:
261                 if self.rp >= 0:
262                     os.close(self.rp)
263                 self.rp = -1
264
265     @property
266     def readable(self):
267         return not self.eof
268     def read(self):
269         with self.lock:
270             try:
271                 data = os.read(self.rp, 1024)
272                 if data == b"":
273                     self.eof = True
274             except IOError:
275                 self.eof = True
276             cbs = list(self.queue)
277             self.queue[:] = []
278         for cb in cbs:
279             cb()
280
281     writable = False
282
283     def call(self, cb):
284         with self.lock:
285             if self.wp < 0:
286                 raise Exception("stopped")
287             self.queue.append(cb)
288             os.write(self.wp, b"a")
289
290     def stop(self):
291         with self.lock:
292             if self.wp >= 0:
293                 os.close(self.wp)
294                 self.wp = -1
295
296 def currentwatcher(io, current):
297     def check(io):
298         if not current:
299             io.stop()
300     io.loopcheck.add(check)