Repository: python-lsp/python-lsp-server
Pull request #499 (2023-12-19 10:49:07)

Make autoimport cache generation non-blocking
On a medium-sized virtual environment, rope takes 10-20 sec to create the module cache. During this time, the language server is unresponsive.

This is not a problem on a local machine where the cache is unlikely emptied, but poses a challenge in cloud environments such as Databricks, where the server runs on a cloud VM which gets pre-emptied regularly.

This PR makes cache initialization non-blocking, hence making the server response immediately after receiving an initialized message.

## How is this tested?

- [x] unit tests
- [x] e2e tests - making sure that the LSP server becomes responsive immediately, and rope working after the cache is created.

tkrabel: @ccordoba12 before diving further into this, I wanted to run the general idea past you.

## What is the problem?

- Currently, `m_initialize` is blocked for 10-20 seconds on rope cache creation.
- While rope creates the cache, it bombards the client with 10k `$/progress` messages before it sends the `ack`. This overwhelms the language client, taking it another 40s until it becomes responsive and starts responding to the `ack` with an `initialized`
- If I make the cache creation non-blocking, but keep sending progress reports at that rate, the client is still overwhelmed

## That's why I suggest the following solution

1. we definitely need to reduce the progress reporting. I deactivated it currently by commenting out the line containing `PylspTaskHandle`, but I'm happy to change that to throttle progress reporting.
2. I made the cache creation non-blocking by running it in another thread, and only sending back completion or code action suggestions back once the thread is done.

CC @bagel897 since she worked on rope a lot.

Any thoughts folks?
bagel897: Is the database consistent during generation? 
You could batch progress reporting or do it per package instead of per module. Regardless cache generation can occupy 100% cpu. You may wanna implement a parallelism limiter. 
ccordoba12: > we definitely need to reduce the progress reporting. I deactivated it currently by commenting out the line containing PylspTaskHandle, but I'm happy to change that to throttle progress reporting.

I'd be in favor of using throttling for this, if possible, to let users know what's happening with autoimport in the meantime.

> I made the cache creation non-blocking by running it in another thread, and only sending back completion or code action suggestions back once the thread is done.

I think this is the right thing to do to avoid blocking the server on `m_initialize`.
tkrabel: @ccordoba12 the PR is ready for review. I throttled the progress reporting to 500ms. Progress begin and progress end messages are unaffected by it. I think this is more than enough :)


Patch:
diff --git a/pylsp/_utils.py b/pylsp/_utils.py
index 9d393b9..f694a36 100644
--- a/pylsp/_utils.py
+++ b/pylsp/_utils.py
@@ -8,6 +8,7 @@ import os
 import pathlib
 import re
 import threading
+import time
 from typing import List, Optional
 
 import docstring_to_markdown
@@ -55,6 +56,23 @@ def debounce(interval_s, keyed_by=None):
     return wrapper
 
 
+def throttle(seconds=1):
+    """Throttles calls to a function evey `seconds` seconds."""
+
+    def decorator(func):
+        @functools.wraps(func)
+        def wrapper(*args, **kwargs):  # pylint: disable=inconsistent-return-statements
+            if not hasattr(wrapper, "last_call"):
+                wrapper.last_call = 0
+            if time.time() - wrapper.last_call >= seconds:
+                wrapper.last_call = time.time()
+                return func(*args, **kwargs)
+
+        return wrapper
+
+    return decorator
+
+
 def find_parents(root, path, names):
     """Find files matching the given names relative to the given path.
 
diff --git a/pylsp/plugins/_rope_task_handle.py b/pylsp/plugins/_rope_task_handle.py
index ed93a05..841d6fe 100644
--- a/pylsp/plugins/_rope_task_handle.py
+++ b/pylsp/plugins/_rope_task_handle.py
@@ -1,11 +1,13 @@
 from __future__ import annotations
 
 import logging
+
 from typing import Callable, ContextManager, List, Optional, Sequence
 
 from rope.base.taskhandle import BaseJobSet, BaseTaskHandle
 
 from pylsp.workspace import Workspace
+from pylsp._utils import throttle
 
 log = logging.getLogger(__name__)
 Report = Callable[[str, int], None]
@@ -55,6 +57,7 @@ class PylspJobSet(BaseJobSet):
         self.count += 1
         self._report()
 
+    @throttle(0.5)
     def _report(self):
         percent = int(self.get_percent_done())
         message = f"{self.job_name} {self.done}/{self.count}"
diff --git a/pylsp/plugins/rope_autoimport.py b/pylsp/plugins/rope_autoimport.py
index ebcc707..b362640 100644
--- a/pylsp/plugins/rope_autoimport.py
+++ b/pylsp/plugins/rope_autoimport.py
@@ -2,6 +2,7 @@
 
 import logging
 from typing import Any, Dict, Generator, List, Optional, Set, Union
+import threading
 
 import parso
 from jedi import Script
@@ -25,6 +26,55 @@ MAX_RESULTS_COMPLETIONS = 1000
 MAX_RESULTS_CODE_ACTIONS = 5
 
 
+class AutoimportCache:
+    """Handles the cache creation."""
+
+    def __init__(self):
+        self.thread = None
+
+    def reload_cache(
+        self,
+        config: Config,
+        workspace: Workspace,
+        files: Optional[List[Document]] = None,
+        single_thread: Optional[bool] = False,
+    ):
+        if self.is_blocked():
+            return
+
+        memory: bool = config.plugin_settings("rope_autoimport").get("memory", False)
+        rope_config = config.settings().get("rope", {})
+        autoimport = workspace._rope_autoimport(rope_config, memory)
+        resources: Optional[List[Resource]] = (
+            None
+            if files is None
+            else [document._rope_resource(rope_config) for document in files]
+        )
+
+        if single_thread:
+            self._reload_cache(workspace, autoimport, resources)
+        else:
+            # Creating the cache may take 10-20s for a environment with 5k python modules. That's
+            # why we decided to move cache creation into its own thread.
+            self.thread = threading.Thread(
+                target=self._reload_cache, args=(workspace, autoimport, resources)
+            )
+            self.thread.start()
+
+    def _reload_cache(
+        self,
+        workspace: Workspace,
+        autoimport: AutoImport,
+        resources: Optional[List[Resource]] = None,
+    ):
+        task_handle = PylspTaskHandle(workspace)
+        autoimport.generate_cache(task_handle=task_handle, resources=resources)
+        autoimport.generate_modules_cache(task_handle=task_handle)
+
+    def is_blocked(self):
+        return self.thread and self.thread.is_alive()
+
+
 @hookimpl
 def pylsp_settings() -> Dict[str, Dict[str, Dict[str, Any]]]:
     # Default rope_completion to disabled
@@ -191,7 +241,7 @@ def pylsp_completions(
         not config.plugin_settings("rope_autoimport")
         .get("completions", {})
         .get("enabled", True)
-    ):
+    ) or cache.is_blocked():
         return []
 
     line = document.lines[position["line"]]
@@ -283,7 +333,7 @@ def pylsp_code_actions(
         not config.plugin_settings("rope_autoimport")
         .get("code_actions", {})
         .get("enabled", True)
-    ):
+    ) or cache.is_blocked():
         return []
 
     log.debug(f"textDocument/codeAction: {document} {range} {context}")
@@ -319,29 +369,13 @@ def pylsp_code_actions(
     return code_actions
 
 
-def _reload_cache(
-    config: Config, workspace: Workspace, files: Optional[List[Document]] = None
-):
-    memory: bool = config.plugin_settings("rope_autoimport").get("memory", False)
-    rope_config = config.settings().get("rope", {})
-    autoimport = workspace._rope_autoimport(rope_config, memory)
-    task_handle = PylspTaskHandle(workspace)
-    resources: Optional[List[Resource]] = (
-        None
-        if files is None
-        else [document._rope_resource(rope_config) for document in files]
-    )
-    autoimport.generate_cache(task_handle=task_handle, resources=resources)
-    autoimport.generate_modules_cache(task_handle=task_handle)
-
-
 @hookimpl
 def pylsp_initialize(config: Config, workspace: Workspace):
     """Initialize AutoImport.
 
     Generates the cache for local and global items.
     """
-    _reload_cache(config, workspace)
+    cache.reload_cache(config, workspace)
 
 
 @hookimpl
@@ -350,13 +384,13 @@ def pylsp_document_did_open(config: Config, workspace: Workspace):
 
     Generates the cache for local and global items.
     """
-    _reload_cache(config, workspace)
+    cache.reload_cache(config, workspace)
 
 
 @hookimpl
 def pylsp_document_did_save(config: Config, workspace: Workspace, document: Document):
     """Update the names associated with this document."""
-    _reload_cache(config, workspace, [document])
+    cache.reload_cache(config, workspace, [document])
 
 
 @hookimpl
@@ -368,6 +402,9 @@ def pylsp_workspace_configuration_changed(config: Config, workspace: Workspace):
     Generates the cache for local and global items.
     """
     if config.plugin_settings("rope_autoimport").get("enabled", False):
-        _reload_cache(config, workspace)
+        cache.reload_cache(config, workspace)
     else:
         log.debug("autoimport: Skipping cache reload.")
+
+
+cache: AutoimportCache = AutoimportCache()


Test patch:
diff --git a/test/plugins/test_autoimport.py b/test/plugins/test_autoimport.py
index 9f02965..6d3eeb5 100644
--- a/test/plugins/test_autoimport.py
+++ b/test/plugins/test_autoimport.py
@@ -15,13 +15,13 @@ from pylsp.config.config import Config
 from pylsp.plugins.rope_autoimport import (
     _get_score,
     _should_insert,
+    cache,
     get_name_or_module,
     get_names,
 )
 from pylsp.plugins.rope_autoimport import (
     pylsp_completions as pylsp_autoimport_completions,
 )
-from pylsp.plugins.rope_autoimport import pylsp_initialize
 from pylsp.workspace import Workspace
 
 
@@ -57,7 +57,7 @@ def autoimport_workspace(tmp_path_factory) -> Workspace:
             }
         }
     )
-    pylsp_initialize(workspace._config, workspace)
+    cache.reload_cache(workspace._config, workspace, single_thread=True)
     yield workspace
     workspace.close()
 
@@ -293,7 +293,6 @@ def test_autoimport_code_actions_and_completions_for_notebook_document(
             }
         },
     )
-
     with patch.object(server._endpoint, "notify") as mock_notify:
         # Expectations:
         # 1. We receive an autoimport suggestion for "os" in the first cell because
@@ -305,13 +304,19 @@ def test_autoimport_code_actions_and_completions_for_notebook_document(
         # 4. We receive an autoimport suggestion for "sys" because it's not already imported.
         # 5. If diagnostics doesn't contain "undefined name ...", we send empty quick fix suggestions.
         send_notebook_did_open(client, ["os", "import os\nos", "os", "sys"])
-        wait_for_condition(lambda: mock_notify.call_count >= 3)
+        wait_for_condition(lambda: mock_notify.call_count >= 4)
+        # We received diagnostics messages for every cell
+        assert all(
+            "textDocument/publishDiagnostics" in c.args
+            for c in mock_notify.call_args_list
+        )
 
     rope_autoimport_settings = server.workspace._config.plugin_settings(
         "rope_autoimport"
     )
     assert rope_autoimport_settings.get("completions", {}).get("enabled", False) is True
     assert rope_autoimport_settings.get("memory", False) is True
+    wait_for_condition(lambda: not cache.thread.is_alive())
 
     # 1.
     quick_fixes = server.code_actions("cell_1_uri", {}, make_context("os", 0, 0, 2))