Catch KeyError within Threadless executors (#1396)

Author: abhinavsingh

.gitignore

@@ -36,6 +36,5 @@ build
 
 pyreverse.png
 profile.svg
-
 *-pre-push
 jaxl-api-credentials*.json

.vscode/settings.json

@@ -41,5 +41,19 @@
   "python.linting.flake8Args": ["--config", ".flake8"],
   "python.linting.mypyEnabled": true,
   "python.formatting.provider": "autopep8",
-  "autoDocstring.docstringFormat": "sphinx"
+  "autoDocstring.docstringFormat": "sphinx",
+  "emeraldwalk.runonsave": {
+    "commands": [
+      {
+        "match": "\\.py$",
+        "isAsync": false,
+        "cmd": "./.venv/bin/autoflake --in-place --remove-all-unused-imports \"${file}\""
+      },
+      {
+        "match": "\\.py$",
+        "isAsync": false,
+        "cmd": "./.venv/bin/isort \"${file}\""
+      }
+    ]
+  }
 }

proxy/core/work/threadless.py

@@ -191,13 +191,21 @@ async def _update_work_events(self, work_id: int) -> None:
             #
             # TODO: Also remove offending work from pool to avoid spin loop.
             elif fileno != -1:
-                self.selector.register(fileno, events=mask, data=work_id)
-                self.registered_events_by_work_ids[work_id][fileno] = mask
-                logger.debug(
-                    'fd#{0} registered for mask#{1} by work#{2}'.format(
-                        fileno, mask, work_id,
-                    ),
-                )
+                try:
+                    self.selector.register(fileno, events=mask, data=work_id)
+                    self.registered_events_by_work_ids[work_id][fileno] = mask
+                    logger.debug(
+                        'fd#{0} registered for mask#{1} by work#{2}'.format(
+                            fileno,
+                            mask,
+                            work_id,
+                        ),
+                    )
+                except KeyError as exc:
+                    logger.warning(
+                        'KeyError when trying to register fd#{0}'.format(fileno),
+                        exc_info=exc,
+                    )
 
     async def _update_conn_pool_events(self) -> None:
         if not self._upstream_conn_pool:

proxy/http/parser/parser.py

@@ -216,6 +216,13 @@ def is_connection_upgrade(self) -> bool:
             self.has_header(b'Connection') and \
             self.has_header(b'Upgrade')
 
+    @property
+    def is_websocket_upgrade(self) -> bool:
+        return (
+            self.is_connection_upgrade
+            and self.header(b'upgrade').lower() == b'websocket'
+        )
+
     @property
     def is_https_tunnel(self) -> bool:
         """Returns true for HTTPS CONNECT tunnel request."""