Skip to content

Commit b70222d

Browse files
committed
perf(lsp): use string.buffer for rpc loop
Avoids some table allocations. In a quick test over 50000 iterations it reduces the time from 130ms to 74 ms For the test setup details see: mfussenegger/nvim-dap#1394 (comment)
1 parent 9d665bd commit b70222d

File tree

1 file changed

+90
-55
lines changed

1 file changed

+90
-55
lines changed

runtime/lua/vim/lsp/rpc.lua

Lines changed: 90 additions & 55 deletions
Original file line numberDiff line numberDiff line change
@@ -49,70 +49,105 @@ local header_start_pattern = ('content'):gsub('%w', function(c)
4949
return '[' .. c .. c:upper() .. ']'
5050
end)
5151

52+
53+
local has_strbuffer, strbuffer = pcall(require, "string.buffer")
54+
5255
--- The actual workhorse.
53-
local function request_parser_loop()
54-
local buffer = '' -- only for header part
55-
while true do
56-
-- A message can only be complete if it has a double CRLF and also the full
57-
-- payload, so first let's check for the CRLFs
58-
local header_end, body_start = buffer:find('\r\n\r\n', 1, true)
59-
-- Start parsing the headers
60-
if header_end then
61-
-- This is a workaround for servers sending initial garbage before
62-
-- sending headers, such as if a bash script sends stdout. It assumes
63-
-- that we know all of the headers ahead of time. At this moment, the
64-
-- only valid headers start with "Content-*", so that's the thing we will
65-
-- be searching for.
66-
-- TODO(ashkan) I'd like to remove this, but it seems permanent :(
67-
local buffer_start = buffer:find(header_start_pattern)
68-
if not buffer_start then
69-
error(
70-
string.format(
71-
"Headers were expected, a different response was received. The server response was '%s'.",
72-
buffer
73-
)
74-
)
75-
end
76-
local header = buffer:sub(buffer_start, header_end + 1)
77-
local content_length = get_content_length(header)
78-
-- Use table instead of just string to buffer the message. It prevents
79-
-- a ton of strings allocating.
80-
-- ref. http://www.lua.org/pil/11.6.html
81-
---@type string[]
82-
local body_chunks = { buffer:sub(body_start + 1) }
83-
local body_length = #body_chunks[1]
84-
-- Keep waiting for data until we have enough.
85-
while body_length < content_length do
86-
---@type string
56+
---@type function
57+
local request_parser_loop
58+
59+
if has_strbuffer then
60+
request_parser_loop = function()
61+
local buf = strbuffer.new()
62+
while true do
63+
local msg = buf:tostring()
64+
local header_end = msg:find('\r\n\r\n', 1, true)
65+
if header_end then
66+
local header = buf:get(header_end + 1)
67+
buf:skip(2) -- skip past header boundary
68+
local content_length = get_content_length(header)
69+
if not content_length then
70+
error("Content-Length not found in headers: " .. header)
71+
end
72+
while #buf < content_length do
73+
local chunk = coroutine.yield()
74+
buf:put(chunk)
75+
end
76+
local body = buf:get(content_length)
77+
local chunk = coroutine.yield(body)
78+
buf:put(chunk)
79+
else
8780
local chunk = coroutine.yield()
88-
or error('Expected more data for the body. The server may have died.') -- TODO hmm.
89-
table.insert(body_chunks, chunk)
90-
body_length = body_length + #chunk
81+
buf:put(chunk)
9182
end
92-
local last_chunk = body_chunks[#body_chunks]
83+
end
84+
end
85+
else
86+
request_parser_loop = function()
87+
local buffer = '' -- only for header part
88+
while true do
89+
-- A message can only be complete if it has a double CRLF and also the full
90+
-- payload, so first let's check for the CRLFs
91+
local header_end, body_start = buffer:find('\r\n\r\n', 1, true)
92+
-- Start parsing the headers
93+
if header_end then
94+
-- This is a workaround for servers sending initial garbage before
95+
-- sending headers, such as if a bash script sends stdout. It assumes
96+
-- that we know all of the headers ahead of time. At this moment, the
97+
-- only valid headers start with "Content-*", so that's the thing we will
98+
-- be searching for.
99+
-- TODO(ashkan) I'd like to remove this, but it seems permanent :(
100+
local buffer_start = buffer:find(header_start_pattern)
101+
if not buffer_start then
102+
error(
103+
string.format(
104+
"Headers were expected, a different response was received. The server response was '%s'.",
105+
buffer
106+
)
107+
)
108+
end
109+
local header = buffer:sub(buffer_start, header_end + 1)
110+
local content_length = get_content_length(header)
111+
-- Use table instead of just string to buffer the message. It prevents
112+
-- a ton of strings allocating.
113+
-- ref. http://www.lua.org/pil/11.6.html
114+
---@type string[]
115+
local body_chunks = { buffer:sub(body_start + 1) }
116+
local body_length = #body_chunks[1]
117+
-- Keep waiting for data until we have enough.
118+
while body_length < content_length do
119+
---@type string
120+
local chunk = coroutine.yield()
121+
or error('Expected more data for the body. The server may have died.') -- TODO hmm.
122+
table.insert(body_chunks, chunk)
123+
body_length = body_length + #chunk
124+
end
125+
local last_chunk = body_chunks[#body_chunks]
93126

94-
body_chunks[#body_chunks] = last_chunk:sub(1, content_length - body_length - 1)
95-
local rest = ''
96-
if body_length > content_length then
97-
rest = last_chunk:sub(content_length - body_length)
98-
end
99-
local body = table.concat(body_chunks)
100-
-- Yield our data.
127+
body_chunks[#body_chunks] = last_chunk:sub(1, content_length - body_length - 1)
128+
local rest = ''
129+
if body_length > content_length then
130+
rest = last_chunk:sub(content_length - body_length)
131+
end
132+
local body = table.concat(body_chunks)
133+
-- Yield our data.
101134

102-
--- @type string
103-
local data = coroutine.yield(body)
104-
or error('Expected more data for the body. The server may have died.')
105-
buffer = rest .. data
106-
else
107-
-- Get more data since we don't have enough.
108-
--- @type string
109-
local data = coroutine.yield()
110-
or error('Expected more data for the header. The server may have died.')
111-
buffer = buffer .. data
135+
--- @type string
136+
local data = coroutine.yield(body)
137+
or error('Expected more data for the body. The server may have died.')
138+
buffer = rest .. data
139+
else
140+
-- Get more data since we don't have enough.
141+
--- @type string
142+
local data = coroutine.yield()
143+
or error('Expected more data for the header. The server may have died.')
144+
buffer = buffer .. data
145+
end
112146
end
113147
end
114148
end
115149

150+
116151
local M = {}
117152

118153
--- Mapping of error codes used by the client

0 commit comments

Comments
 (0)