python: Properly abort when over request limit.
[ashd.git] / python / ashd-wsgi
index af2c332..01c158f 100755 (executable)
@@ -1,17 +1,18 @@
 #!/usr/bin/python
 
-import sys, os, getopt, threading, time
-import ashd.proto, ashd.util, ashd.perf
+import sys, os, getopt, threading, logging, time
+import ashd.proto, ashd.util, ashd.perf, ashd.serve
 try:
     import pdm.srv
 except:
     pdm = None
 
 def usage(out):
-    out.write("usage: ashd-wsgi [-hA] [-m PDM-SPEC] [-p MODPATH] [-l REQLIMIT] HANDLER-MODULE [ARGS...]\n")
+    out.write("usage: ashd-wsgi [-hAL] [-m PDM-SPEC] [-p MODPATH] [-l REQLIMIT] HANDLER-MODULE [ARGS...]\n")
 
 reqlimit = 0
 modwsgi_compat = False
+setlog = True
 opts, args = getopt.getopt(sys.argv[1:], "+hAp:l:m:")
 for o, a in opts:
     if o == "-h":
@@ -19,6 +20,8 @@ for o, a in opts:
         sys.exit(0)
     elif o == "-p":
         sys.path.insert(0, a)
+    elif o == "-L":
+        setlog = False
     elif o == "-A":
         modwsgi_compat = True
     elif o == "-l":
@@ -29,6 +32,9 @@ for o, a in opts:
 if len(args) < 1:
     usage(sys.stderr)
     sys.exit(1)
+if setlog:
+    logging.basicConfig(format="ashd-wsgi(%(name)s): %(levelname)s: %(message)s")
+log = logging.getLogger("ashd-wsgi")
 
 try:
     handlermod = __import__(args[0], fromlist = ["dummy"])
@@ -46,10 +52,6 @@ else:
         sys.exit(1)
     handler = handlermod.application
 
-class closed(IOError):
-    def __init__(self):
-        super(closed, self).__init__("The client has closed the connection.")
-
 cwd = os.getcwd()
 def absolutify(path):
     if path[0] != '/':
@@ -89,7 +91,7 @@ def unquoteurl(url):
             buf += c
     return buf
 
-def dowsgi(req):
+def mkenv(req):
     env = {}
     env["wsgi.version"] = 1, 0
     for key, val in req.headers:
@@ -121,11 +123,19 @@ def dowsgi(req):
     env["SCRIPT_NAME"] = name
     env["PATH_INFO"] = pi
     if "Host" in req: env["SERVER_NAME"] = req["Host"]
+    if "X-Ash-Server-Address" in req: env["SERVER_ADDR"] = req["X-Ash-Server-Address"]
     if "X-Ash-Server-Port" in req: env["SERVER_PORT"] = req["X-Ash-Server-Port"]
     if "X-Ash-Protocol" in req and req["X-Ash-Protocol"] == "https": env["HTTPS"] = "on"
     if "X-Ash-Address" in req: env["REMOTE_ADDR"] = req["X-Ash-Address"]
-    if "Content-Type" in req: env["CONTENT_TYPE"] = req["Content-Type"]
-    if "Content-Length" in req: env["CONTENT_LENGTH"] = req["Content-Length"]
+    if "X-Ash-Port" in req: env["REMOTE_PORT"] = req["X-Ash-Port"]
+    if "Content-Type" in req:
+        env["CONTENT_TYPE"] = req["Content-Type"]
+        # The CGI specification does not strictly require this, but
+        # many actualy programs and libraries seem to.
+        del env["HTTP_CONTENT_TYPE"]
+    if "Content-Length" in req:
+        env["CONTENT_LENGTH"] = req["Content-Length"]
+        del env["HTTP_CONTENT_LENGTH"]
     if "X-Ash-File" in req: env["SCRIPT_FILENAME"] = absolutify(req["X-Ash-File"])
     if "X-Ash-Protocol" in req: env["wsgi.url_scheme"] = req["X-Ash-Protocol"]
     env["wsgi.input"] = req.sk
@@ -133,100 +143,54 @@ def dowsgi(req):
     env["wsgi.multithread"] = True
     env["wsgi.multiprocess"] = False
     env["wsgi.run_once"] = False
+    return env
 
-    resp = []
-    respsent = []
+if reqlimit != 0:
+    guard = ashd.serve.abortlimiter(reqlimit).call
+else:
+    guard = lambda fun: fun()
 
-    def flushreq():
-        if not respsent:
-            if not resp:
-                raise Exception, "Trying to write data before starting response."
-            status, headers = resp
-            respsent[:] = [True]
-            try:
-                req.sk.write("HTTP/1.1 %s\n" % status)
-                for nm, val in headers:
-                    req.sk.write("%s: %s\n" % (nm, val))
-                req.sk.write("\n")
-            except IOError:
-                raise closed()
+class reqthread(ashd.serve.wsgithread):
+    def __init__(self, req):
+        super(reqthread, self).__init__()
+        self.req = req.dup()
+    
+    def handlewsgi(self):
+        return handler(self.env, self.startreq)
 
-    def write(data):
-        if not data:
-            return
-        flushreq()
+    def writehead(self, status, headers):
         try:
-            req.sk.write(data)
-            req.sk.flush()
+            self.req.sk.write("HTTP/1.1 %s\n" % status)
+            for nm, val in headers:
+                self.req.sk.write("%s: %s\n" % (nm, val))
+            self.req.sk.write("\n")
         except IOError:
-            raise closed()
+            raise ashd.serve.closed()
 
-    def startreq(status, headers, exc_info = None):
-        if resp:
-            if exc_info:                # Interesting, this...
-                try:
-                    if respsent:
-                        raise exc_info[0], exc_info[1], exc_info[2]
-                finally:
-                    exc_info = None     # CPython GC bug?
-            else:
-                raise Exception, "Can only start responding once."
-        resp[:] = status, headers
-        return write
+    def writedata(self, data):
+        try:
+            self.req.sk.write(data)
+            self.req.sk.flush()
+        except IOError:
+            raise ashd.serve.closed()
 
-    reqevent = ashd.perf.request(env)
-    exc = (None, None, None)
-    try:
-        respiter = handler(env, startreq)
+    def handle(self):
+        self.env = mkenv(self.req)
+        reqevent = ashd.perf.request(self.env)
+        exc = (None, None, None)
         try:
-            try:
-                for data in respiter:
-                    write(data)
-                if resp:
-                    flushreq()
-            except closed:
-                pass
+            super(reqthread, self).handle()
+            if self.status:
+                reqevent.response([self.status, self.headers])
+        except:
+            exc = sys.exc_info()
+            raise
         finally:
-            if hasattr(respiter, "close"):
-                respiter.close()
-        if resp:
-            reqevent.response(resp)
-    except:
-        exc = sys.exc_info()
-    finally:
-        reqevent.__exit__(*exc)
-
-flightlock = threading.Condition()
-inflight = 0
+            reqevent.__exit__(*exc)
 
-class reqthread(threading.Thread):
-    def __init__(self, req):
-        super(reqthread, self).__init__(name = "Request handler")
-        self.req = req.dup()
-    
     def run(self):
-        global inflight
         try:
-            flightlock.acquire()
-            try:
-                if reqlimit != 0:
-                    start = time.time()
-                    while inflight >= reqlimit:
-                        flightlock.wait(10)
-                        if time.time() - start > 10:
-                            os.abort()
-                inflight += 1
-            finally:
-                flightlock.release()
-            try:
-                dowsgi(self.req)
-            finally:
-                flightlock.acquire()
-                try:
-                    inflight -= 1
-                    flightlock.notify()
-                finally:
-                    flightlock.release()
+            guard(super(reqthread, self).run)
         finally:
             self.req.close()