Skip to content

Commit 8413427

Browse files
mfusseneggermfussenegger
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 2fdb688 commit 8413427

File tree

1 file changed

+87
-57
lines changed

1 file changed

+87
-57
lines changed

‎runtime/lua/vim/lsp/rpc.lua‎

Lines changed: 87 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -25,8 +25,8 @@ local function get_content_length(header)
2525
if line == '' then
2626
break
2727
end
28-
local key, value = line:match('^%s*(%S+)%s*:%s*(.+)%s*$')
29-
if key:lower() == 'content-length' then
28+
local key, value = line:match('^%s*(%S+)%s*:%s*(%d+)%s*$')
29+
if key and key:lower() == 'content-length' then
3030
return tonumber(value)
3131
end
3232
end
@@ -39,66 +39,96 @@ local header_start_pattern = ('content'):gsub('%w', function(c)
3939
return '[' .. c .. c:upper() .. ']'
4040
end)
4141

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

84-
body_chunks[#body_chunks] = last_chunk:sub(1, content_length - body_length - 1)
85-
local rest = ''
86-
if body_length > content_length then
87-
rest = last_chunk:sub(content_length - body_length)
88-
end
89-
local body = table.concat(body_chunks)
90-
-- Yield our data.
113+
body_chunks[#body_chunks] = last_chunk:sub(1, content_length - body_length - 1)
114+
local rest = ''
115+
if body_length > content_length then
116+
rest = last_chunk:sub(content_length - body_length)
117+
end
118+
local body = table.concat(body_chunks)
119+
-- Yield our data.
91120

92-
--- @type string
93-
local data = coroutine.yield(body)
94-
or error('Expected more data for the body. The server may have died.')
95-
buffer = rest .. data
96-
else
97-
-- Get more data since we don't have enough.
98-
--- @type string
99-
local data = coroutine.yield()
100-
or error('Expected more data for the header. The server may have died.')
101-
buffer = buffer .. data
121+
--- @type string
122+
local data = coroutine.yield(body)
123+
or error('Expected more data for the body. The server may have died.')
124+
buffer = rest .. data
125+
else
126+
-- Get more data since we don't have enough.
127+
--- @type string
128+
local data = coroutine.yield()
129+
or error('Expected more data for the header. The server may have died.')
130+
buffer = buffer .. data
131+
end
102132
end
103133
end
104134
end

0 commit comments

Comments
 (0)