feat: add semaphore to limit concurrency in capability requests (#23)

* feat: add semaphore to limit concurrency in capability requests

* refactor: expand semaphore scope to cover entire processing blocks

Author: observerw

src/lsap/capability/definition.py

@@ -4,8 +4,9 @@
 from functools import cached_property
 from typing import override
 
+import anyio
 import asyncer
-from attrs import define
+from attrs import define, field
 from lsp_client.capability.request import (
     WithRequestDeclaration,
     WithRequestDefinition,
@@ -24,6 +25,8 @@
 
 @define
 class DefinitionCapability(Capability[DefinitionRequest, DefinitionResponse]):
+    resolve_sem: anyio.Semaphore = field(default=anyio.Semaphore(32))
+
     @cached_property
     def locate(self) -> LocateCapability:
         return LocateCapability(self.client)
@@ -69,12 +72,13 @@ async def __call__(self, req: DefinitionRequest) -> DefinitionResponse | None:
         async with asyncer.create_task_group() as tg:
 
             async def resolve_item(loc: Location) -> SymbolCodeInfo | None:
-                target_file_path = self.client.from_uri(loc.uri)
-                if symbol_info := await self.symbol.resolve(
-                    target_file_path,
-                    loc.range.start,
-                ):
-                    return symbol_info
+                async with self.resolve_sem:
+                    target_file_path = self.client.from_uri(loc.uri)
+                    if symbol_info := await self.symbol.resolve(
+                        target_file_path,
+                        loc.range.start,
+                    ):
+                        return symbol_info
                 return None
 
             infos = [tg.soonify(resolve_item)(loc) for loc in locations]

src/lsap/capability/outline.py

@@ -4,8 +4,9 @@
 from pathlib import Path
 from typing import override
 
+import anyio
 import asyncer
-from attrs import define
+from attrs import define, field
 from lsp_client.capability.request import WithRequestDocumentSymbol, WithRequestHover
 from lsprotocol.types import DocumentSymbol
 from lsprotocol.types import Position as LSPPosition
@@ -22,6 +23,8 @@
 
 @define
 class OutlineCapability(Capability[OutlineRequest, OutlineResponse]):
+    hover_sem: anyio.Semaphore = field(default=anyio.Semaphore(32))
+
     @override
     async def __call__(self, req: OutlineRequest) -> OutlineResponse | None:
         symbols = await ensure_capability(
@@ -117,7 +120,8 @@ def _make_item(
         )
 
     async def _fill_hover(self, item: SymbolDetailInfo, pos: LSPPosition) -> None:
-        if hover := await ensure_capability(
-            self.client, WithRequestHover
-        ).request_hover(item.file_path, pos):
-            item.hover = clean_hover_content(hover.value)
+        async with self.hover_sem:
+            if hover := await ensure_capability(
+                self.client, WithRequestHover
+            ).request_hover(item.file_path, pos):
+                item.hover = clean_hover_content(hover.value)

src/lsap/capability/reference.py

@@ -1,7 +1,8 @@
 from functools import cached_property
 
+import anyio
 import asyncer
-from attrs import Factory, define
+from attrs import Factory, define, field
 from lsp_client.capability.request import (
     WithRequestDocumentSymbol,
     WithRequestHover,
@@ -29,6 +30,7 @@
 @define
 class ReferenceCapability(Capability[ReferenceRequest, ReferenceResponse]):
     _cache: PaginationCache[ReferenceItem] = Factory(PaginationCache)
+    process_sem: anyio.Semaphore = field(default=anyio.Semaphore(32))
 
     @cached_property
     def locate(self) -> LocateCapability:
@@ -92,56 +94,57 @@ async def _process_reference(
         context_lines: int,
         items: list[ReferenceItem],
     ) -> None:
-        file_path = self.client.from_uri(loc.uri)
-        content = await self.client.read_file(file_path)
-        reader = DocumentReader(content)
-
-        range = loc.range
-        context_range = LSPRange(
-            start=LSPPosition(
-                line=max(0, range.start.line - context_lines), character=0
-            ),
-            end=LSPPosition(line=range.end.line + context_lines + 1, character=0),
-        )
-        if not (snippet := reader.read(context_range, trim_empty=True)):
-            return
-
-        symbol: SymbolDetailInfo | None = None
-        if (
-            symbols := await ensure_capability(
-                self.client, WithRequestDocumentSymbol
-            ).request_document_symbol_list(file_path)
-        ) and (match := symbol_at(symbols, range.start)):
-            path, sym = match
-            kind = SymbolKind.from_lsp(sym.kind)
-
-            symbol = SymbolDetailInfo(
-                file_path=file_path,
-                name=sym.name,
-                path=path,
-                kind=kind,
-                detail=sym.detail,
-                range=Range(
-                    start=Position.from_lsp(sym.range.start),
-                    end=Position.from_lsp(sym.range.end),
+        async with self.process_sem:
+            file_path = self.client.from_uri(loc.uri)
+            content = await self.client.read_file(file_path)
+            reader = DocumentReader(content)
+
+            range = loc.range
+            context_range = LSPRange(
+                start=LSPPosition(
+                    line=max(0, range.start.line - context_lines), character=0
                 ),
+                end=LSPPosition(line=range.end.line + context_lines + 1, character=0),
             )
-
-            if hover := await ensure_capability(
-                self.client, WithRequestHover
-            ).request_hover(file_path, range.start):
-                symbol.hover = clean_hover_content(hover.value)
-
-        items.append(
-            ReferenceItem(
-                location=LSAPLocation(
+            if not (snippet := reader.read(context_range, trim_empty=True)):
+                return
+
+            symbol: SymbolDetailInfo | None = None
+            if (
+                symbols := await ensure_capability(
+                    self.client, WithRequestDocumentSymbol
+                ).request_document_symbol_list(file_path)
+            ) and (match := symbol_at(symbols, range.start)):
+                path, sym = match
+                kind = SymbolKind.from_lsp(sym.kind)
+
+                symbol = SymbolDetailInfo(
                     file_path=file_path,
+                    name=sym.name,
+                    path=path,
+                    kind=kind,
+                    detail=sym.detail,
                     range=Range(
-                        start=Position.from_lsp(range.start),
-                        end=Position.from_lsp(range.end),
+                        start=Position.from_lsp(sym.range.start),
+                        end=Position.from_lsp(sym.range.end),
                     ),
-                ),
-                code=snippet.content,
-                symbol=symbol,
+                )
+
+                if hover := await ensure_capability(
+                    self.client, WithRequestHover
+                ).request_hover(file_path, range.start):
+                    symbol.hover = clean_hover_content(hover.value)
+
+            items.append(
+                ReferenceItem(
+                    location=LSAPLocation(
+                        file_path=file_path,
+                        range=Range(
+                            start=Position.from_lsp(range.start),
+                            end=Position.from_lsp(range.end),
+                        ),
+                    ),
+                    code=snippet.content,
+                    symbol=symbol,
+                )
             )
-        )

src/lsap/capability/rename.py

@@ -6,8 +6,9 @@
 from pathlib import Path
 from typing import override
 
+import anyio
 import asyncer
-from attrs import define
+from attrs import define, field
 from lsp_client.capability.request import WithRequestRename
 from lsp_client.protocol import CapabilityClientProtocol
 from lsp_client.utils.types import lsp_type
@@ -141,6 +142,8 @@ def should_exclude(uri: str) -> bool:
 
 @define
 class RenamePreviewCapability(Capability[RenamePreviewRequest, RenamePreviewResponse]):
+    file_sem: anyio.Semaphore = field(default=anyio.Semaphore(32))
+
     @cached_property
     def locate(self) -> LocateCapability:
         return LocateCapability(self.client)
@@ -208,44 +211,51 @@ async def _to_file_change(
         *,
         reader: DocumentReader | None = None,
     ) -> RenameFileChange | None:
-        if reader is None:
-            content = await self.client.read_file(
-                self.client.from_uri(uri, relative=False)
-            )
-            reader = DocumentReader(content)
-
-        diffs: list[RenameDiff] = []
-        for edit in edits:
-            start, end = edit.range.start, edit.range.end
-            line_raw = reader.get_line(start.line, keepends=True)
-            if line_raw is None:
-                continue
-
-            original_line = line_raw.rstrip("\r\n")
-            new_text = get_edit_text(edit)
-
-            if start.line == end.line:
-                modified_line = (
-                    line_raw[: start.character] + new_text + line_raw[end.character :]
-                ).rstrip("\r\n")
-            else:
-                modified_line = new_text
-
-            diffs.append(
-                RenameDiff(
-                    line=start.line + 1,
-                    original=original_line,
-                    modified=modified_line,
+        async with self.file_sem:
+            if reader is None:
+                content = await self.client.read_file(
+                    self.client.from_uri(uri, relative=False)
+                )
+                reader = DocumentReader(content)
+
+            diffs: list[RenameDiff] = []
+            for edit in edits:
+                start, end = edit.range.start, edit.range.end
+                line_raw = reader.get_line(start.line, keepends=True)
+                if line_raw is None:
+                    continue
+
+                original_line = line_raw.rstrip("\r\n")
+                new_text = get_edit_text(edit)
+
+                if start.line == end.line:
+                    modified_line = (
+                        line_raw[: start.character]
+                        + new_text
+                        + line_raw[end.character :]
+                    ).rstrip("\r\n")
+                else:
+                    modified_line = new_text
+
+                diffs.append(
+                    RenameDiff(
+                        line=start.line + 1,
+                        original=original_line,
+                        modified=modified_line,
+                    )
                 )
-            )
 
-        if diffs:
-            return RenameFileChange(file_path=self.client.from_uri(uri), diffs=diffs)
-        return None
+            if diffs:
+                return RenameFileChange(
+                    file_path=self.client.from_uri(uri), diffs=diffs
+                )
+            return None
 
 
 @define
 class RenameExecuteCapability(Capability[RenameExecuteRequest, RenameExecuteResponse]):
+    file_sem: anyio.Semaphore = field(default=anyio.Semaphore(32))
+
     @override
     async def __call__(self, req: RenameExecuteRequest) -> RenameExecuteResponse | None:
         cached = _preview_cache.get(req.rename_id)
@@ -301,37 +311,42 @@ async def _to_file_change(
         *,
         reader: DocumentReader | None = None,
     ) -> RenameFileChange | None:
-        if reader is None:
-            content = await self.client.read_file(
-                self.client.from_uri(uri, relative=False)
-            )
-            reader = DocumentReader(content)
-
-        diffs: list[RenameDiff] = []
-        for edit in edits:
-            start, end = edit.range.start, edit.range.end
-            line_raw = reader.get_line(start.line, keepends=True)
-            if line_raw is None:
-                continue
-
-            original_line = line_raw.rstrip("\r\n")
-            new_text = get_edit_text(edit)
-
-            if start.line == end.line:
-                modified_line = (
-                    line_raw[: start.character] + new_text + line_raw[end.character :]
-                ).rstrip("\r\n")
-            else:
-                modified_line = new_text
-
-            diffs.append(
-                RenameDiff(
-                    line=start.line + 1,
-                    original=original_line,
-                    modified=modified_line,
+        async with self.file_sem:
+            if reader is None:
+                content = await self.client.read_file(
+                    self.client.from_uri(uri, relative=False)
+                )
+                reader = DocumentReader(content)
+
+            diffs: list[RenameDiff] = []
+            for edit in edits:
+                start, end = edit.range.start, edit.range.end
+                line_raw = reader.get_line(start.line, keepends=True)
+                if line_raw is None:
+                    continue
+
+                original_line = line_raw.rstrip("\r\n")
+                new_text = get_edit_text(edit)
+
+                if start.line == end.line:
+                    modified_line = (
+                        line_raw[: start.character]
+                        + new_text
+                        + line_raw[end.character :]
+                    ).rstrip("\r\n")
+                else:
+                    modified_line = new_text
+
+                diffs.append(
+                    RenameDiff(
+                        line=start.line + 1,
+                        original=original_line,
+                        modified=modified_line,
+                    )
                 )
-            )
 
-        if diffs:
-            return RenameFileChange(file_path=self.client.from_uri(uri), diffs=diffs)
-        return None
+            if diffs:
+                return RenameFileChange(
+                    file_path=self.client.from_uri(uri), diffs=diffs
+                )
+            return None