File size: 7,031 Bytes
d1ceb73 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 |
""" Language Server stdio-mode readers
Parts of this code are derived from:
> https://github.com/palantir/python-jsonrpc-server/blob/0.2.0/pyls_jsonrpc/streams.py#L83 # noqa
> https://github.com/palantir/python-jsonrpc-server/blob/45ed1931e4b2e5100cc61b3992c16d6f68af2e80/pyls_jsonrpc/streams.py # noqa
> > MIT License https://github.com/palantir/python-jsonrpc-server/blob/0.2.0/LICENSE
> > Copyright 2018 Palantir Technologies, Inc.
"""
# pylint: disable=broad-except
import asyncio
import io
import os
from concurrent.futures import ThreadPoolExecutor
from typing import List, Optional, Text
from tornado.concurrent import run_on_executor
from tornado.gen import convert_yielded
from tornado.httputil import HTTPHeaders
from tornado.ioloop import IOLoop
from tornado.queues import Queue
from traitlets import Float, Instance, default
from traitlets.config import LoggingConfigurable
from .non_blocking import make_non_blocking
class LspStdIoBase(LoggingConfigurable):
"""Non-blocking, queued base for communicating with stdio Language Servers"""
executor = None
stream = Instance( # type:ignore[assignment]
io.RawIOBase, help="the stream to read/write"
) # type: io.RawIOBase
queue = Instance(Queue, help="queue to get/put")
def __repr__(self): # pragma: no cover
return "<{}(parent={})>".format(self.__class__.__name__, self.parent)
def __init__(self, **kwargs):
super().__init__(**kwargs)
self.log.debug("%s initialized", self)
self.executor = ThreadPoolExecutor(max_workers=1)
def close(self):
self.stream.close()
self.log.debug("%s closed", self)
class LspStdIoReader(LspStdIoBase):
"""Language Server stdio Reader
Because non-blocking (but still synchronous) IO is used, rudimentary
exponential backoff is used.
"""
max_wait = Float(help="maximum time to wait on idle stream").tag(config=True)
min_wait = Float(0.05, help="minimum time to wait on idle stream").tag(config=True)
next_wait = Float(0.05, help="next time to wait on idle stream").tag(config=True)
@default("max_wait")
def _default_max_wait(self):
return 0.1 if os.name == "nt" else self.min_wait * 2
async def sleep(self):
"""Simple exponential backoff for sleeping"""
if self.stream.closed: # pragma: no cover
return
self.next_wait = min(self.next_wait * 2, self.max_wait)
try:
await asyncio.sleep(self.next_wait)
except Exception: # pragma: no cover
pass
def wake(self):
"""Reset the wait time"""
self.wait = self.min_wait
async def read(self) -> None:
"""Read from a Language Server until it is closed"""
make_non_blocking(self.stream)
while not self.stream.closed:
message = None
try:
message = await self.read_one()
if not message:
await self.sleep()
continue
else:
self.wake()
IOLoop.current().add_callback(self.queue.put_nowait, message)
except Exception as e: # pragma: no cover
self.log.exception(
"%s couldn't enqueue message: %s (%s)", self, message, e
)
await self.sleep()
async def _read_content(
self, length: int, max_parts=1000, max_empties=200
) -> Optional[bytes]:
"""Read the full length of the message unless exceeding max_parts or
max_empties empty reads occur.
See https://github.com/jupyter-lsp/jupyterlab-lsp/issues/450
Crucial docs or read():
"If the argument is positive, and the underlying raw
stream is not interactive, multiple raw reads may be issued
to satisfy the byte count (unless EOF is reached first)"
Args:
- length: the content length
- max_parts: prevent absurdly long messages (1000 parts is several MBs):
1 part is usually sufficient but not enough for some long
messages 2 or 3 parts are often needed.
"""
raw = None
raw_parts: List[bytes] = []
received_size = 0
while received_size < length and len(raw_parts) < max_parts and max_empties > 0:
part = None
try:
part = self.stream.read(length - received_size)
except OSError: # pragma: no cover
pass
if part is None:
max_empties -= 1
await self.sleep()
continue
received_size += len(part)
raw_parts.append(part)
if raw_parts:
raw = b"".join(raw_parts)
if len(raw) != length: # pragma: no cover
self.log.warning(
f"Readout and content-length mismatch: {len(raw)} vs {length};"
f"remaining empties: {max_empties}; remaining parts: {max_parts}"
)
return raw
async def read_one(self) -> Text:
"""Read a single message"""
message = ""
headers = HTTPHeaders()
line = await convert_yielded(self._readline())
if line:
while line and line.strip():
headers.parse_line(line)
line = await convert_yielded(self._readline())
content_length = int(headers.get("content-length", "0"))
if content_length:
raw = await self._read_content(length=content_length)
if raw is not None:
message = raw.decode("utf-8").strip()
else: # pragma: no cover
self.log.warning(
"%s failed to read message of length %s",
self,
content_length,
)
return message
@run_on_executor
def _readline(self) -> Text:
"""Read a line (or immediately return None)"""
try:
return self.stream.readline().decode("utf-8").strip()
except OSError: # pragma: no cover
return ""
class LspStdIoWriter(LspStdIoBase):
"""Language Server stdio Writer"""
async def write(self) -> None:
"""Write to a Language Server until it closes"""
while not self.stream.closed:
message = await self.queue.get()
try:
body = message.encode("utf-8")
response = "Content-Length: {}\r\n\r\n{}".format(len(body), message)
await convert_yielded(self._write_one(response.encode("utf-8")))
except Exception: # pragma: no cover
self.log.exception("%s couldn't write message: %s", self, response)
finally:
self.queue.task_done()
@run_on_executor
def _write_one(self, message) -> None:
self.stream.write(message)
self.stream.flush()
|