Add support for non-blocking HTTP handlers

Author: rschlaikjer

python/mujinasync/asynchttp.py

@@ -1,13 +1,26 @@
 # -*- coding: utf-8 -*-
 
+from typing import Optional, Any
+
 from .asynctcp import TcpServer, TcpConnection, TcpClient
 
 import logging
 log = logging.getLogger(__name__)
 
+class HttpResponseProcessing(Exception):
+    # Exception class that user logic can throw during HandleHttpRequest to indicate that they have received this
+    # request, but are not yet ready to return the response (e.g, calling a slow external resource), and the
+    # connection should not block for a response.
+    # If this is thrown, then the active request will be stashed on the connection and the http server will
+    # continue processing any other active requests
+    pass
+
 
 class HttpConnection(TcpConnection):
-    pass
+    # If this connection has already parsed a request, and that request has been pushed to the caller but backpressured,
+    # then it will be stashed on the connection so that it can be re-submitted each spin until the user logic is ready
+    # to return a response for it.
+    pendingRequest: Optional["HttpRequest"]  = None
 
 class HttpRequest(object):
     method = 'GET'
@@ -19,6 +32,10 @@ class HttpRequest(object):
 
     response = None
 
+    # For user programs that make use of non-blocking handlers
+    # Any information necessary to track the progress of the request can be stored here
+    userState: Optional[Any] = None
+
     def __repr__(self):
         return '<%s(%s)>' % (self.__class__.__name__, ', '.join([
             ('%s=%r' % (key, getattr(self, key)))
@@ -73,19 +90,34 @@ def _HandleHttpReceiveOnce(self, connection):
         response = None
         try:
             response = self._HandleHttpRequest(connection, request)
+        except HttpResponseProcessing:
+            # If the current request would block, then cache it and re-submit it next loop to check if it is complete.
+            # Send no response on the connection here - the request is still processing
+            connection.pendingRequest = request
+            connection.hasPendingWork = True
+            return
         except Exception as e:
             log.exception('caught exception while handling http request %r: %s', request, e)
-        finally:
-            if response is None:
-                response = HttpResponse(request, statusCode=500, statusText='Internal Server Error')
-            log.verbose('sending http response: %r', response)
-            self._SendHttpResponse(connection, request, response)
+
+        if response is None:
+            response = HttpResponse(request, statusCode=500, statusText='Internal Server Error')
+
+        log.verbose('sending http response: %r', response)
+        self._SendHttpResponse(connection, request, response)
+
         return True # handled one request, try next one
 
     def _HandleHttpRequest(self, connection, request):
         return self._CallApi('HandleHttpRequest', request=request, connection=connection, server=self)
 
     def _TryReceiveHttpRequest(self, connection):
+        # If this connection already has a pending request, pop it off and return it
+        if connection.pendingRequest:
+            request = connection.pendingRequest
+            connection.pendingRequest = None
+            connection.hasPendingWork = False
+            return request
+
         bufferData = connection.receiveBuffer.readView.tobytes()
         if b'\r\n\r\n' not in bufferData:
             if len(bufferData) > 10240:

python/mujinasync/asynctcp.py

@@ -70,13 +70,15 @@ class TcpConnection(object):
     closeType = None # Immediate, AfterSend
     sendBuffer = None # buffer to hold data waiting to be sent
     receiveBuffer = None # buffer to hold data received before consumption
+    hasPendingWork: bool # Should this socket be submitted as a 'readable' socket even if no new data is received?
 
     def __init__(self, connectionSocket, remoteAddress):
         self.connectionSocket = connectionSocket
         self.remoteAddress = remoteAddress
         self.closeType = None
         self.sendBuffer = TcpBuffer()
         self.receiveBuffer = TcpBuffer()
+        self.hasPendingWork = False
 
     def __repr__(self):
         return '<%s(%s)>' % (self.__class__.__name__, ', '.join([
@@ -295,7 +297,7 @@ def SpinOnce(self, timeout=0):
         :param timeout: in seconds, pass in 0 to not wait for socket events, otherwise, will wait up to specified timeout
         """
         newConnections = [] # list of tuple (serverClient, connection)
-        
+
         # construct a list of connections to select on
         rsockets = []
         wsockets = []
@@ -330,7 +332,7 @@ def SpinOnce(self, timeout=0):
                 client._connections.append(connection)
                 newConnections.append((client, connection))
                 timeout = 0 # force no wait at select later since we have a new connection to report right away
-        
+
         # pool all the sockets
         socketConnections = {}
         for serverClient in self._servers + self._clients:
@@ -436,6 +438,12 @@ def SpinOnce(self, timeout=0):
                 connection.connectionSocket = None
             serverClient._connections.remove(connection)
 
+        # Handle server sockets that are processing non-blocking work
+        for server in self._servers:
+            for connection in server._connections:
+                if connection.hasPendingWork:
+                    receivedConnections.append((server, connection))
+
         # let user code run at the very end
         for serverClient, connection in newConnections:
             serverClient._HandleTcpConnect(connection)