Merge #39000 MessageStream abstraction

runtime/lua/vim/lsp/rpc.lua

@@ -1,6 +1,6 @@
 local log = require('vim.lsp.log')
 local protocol = require('vim.lsp.protocol')
-local lsp_transport = require('vim.lsp._transport')
+local vim_transport = require('vim.net._transport')
 local strbuffer = require('vim._core.stringbuffer')
 local validate = vim.validate
 
@@ -25,16 +25,18 @@ end
 --- We ignore lines ending with `\n` that don't contain `content-length`, since some servers
 --- write log to standard output and there's no way to avoid it.
 --- See https://github.com/neovim/neovim/pull/35743#pullrequestreview-3379705828
---- @param header string The header to parse
+--- @param ptr vim._core.stringbuffer.ptr The ptr to buffer to parse
+--- @param start integer The starting index of the buffer to parse, 0-based
+--- @param len integer The length of the header to parse
 --- @return integer
-local function get_content_length(header)
+local function get_content_length(ptr, start, len)
   local state = 'name'
-  local i, len = 1, #header
+  local i, end_ = start, start + len
   local j, name = 1, 'content-length'
   local buf = strbuffer.new()
   local digit = true
-  while i <= len do
-    local c = header:byte(i)
+  while i < end_ do
+    local c = ptr[i]
     if state == 'name' then
       if c >= 65 and c <= 90 then -- lower case
         c = c + 32
@@ -54,7 +56,7 @@ local function get_content_length(header)
         i = i - 1
       end
     elseif state == 'value' then
-      if c == 13 and header:byte(i + 1) == 10 then -- must end with \r\n
+      if c == 13 and ptr[i + 1] == 10 then -- must end with \r\n
         local value = buf:get()
         if digit then
           return vim._assert_integer(value)
@@ -73,7 +75,11 @@ local function get_content_length(header)
     end
     i = i + 1
   end
-  error('Content-Length not found in header: ' .. header)
+  local header = strbuffer.new()
+  for k = start, end_ - 1 do
+    header:put(string.char(ptr[k]))
+  end
+  error('Content-Length not found in header: ' .. header:tostring())
 end
 
 local M = {}
@@ -190,33 +196,36 @@ local default_dispatchers = {
   end,
 }
 
---- @async
-local function message_decoder()
-  local strbuf = strbuffer.new()
-  while true do
-    local header_len ---@type integer?
-    local ptr, len = strbuf:ref()
-    for i = 0, len - 4 do
-      -- Find the header boundary "\r\n\r\n"
-      -- (compare bytes instead of string.find(), to avoid a string alloc).
-      if ptr[i] == 13 and ptr[i + 1] == 10 and ptr[i + 2] == 13 and ptr[i + 3] == 10 then
-        header_len = i + 2
-        break
-      end
-    end
-    if header_len then
-      local header = strbuf:get(header_len)
-      strbuf:skip(2) -- skip past header boundary
-      local content_length = get_content_length(header)
-      while strbuffer.len(strbuf) < content_length do
-        strbuf:put(coroutine.yield())
-      end
-      local body = strbuf:get(content_length)
-      strbuf:put(coroutine.yield(body))
-    else
-      strbuf:put(coroutine.yield())
+--- Parse one `Content-Length` framed message from `strbuf`.
+---
+--- Returns a body after consuming one full frame, returns nil if more bytes are needed.
+--- Raises an error if the buffered data is not a valid frame.
+---
+---@param strbuf vim._core.stringbuffer
+---@return string?
+local function message_decoder(strbuf)
+  local header_len ---@type integer?
+  local ptr, len = strbuf:ref()
+  for i = 0, len - 4 do
+    -- Find the header boundary "\r\n\r\n"
+    -- (compare bytes instead of string.find(), to avoid a string alloc).
+    if ptr[i] == 13 and ptr[i + 1] == 10 and ptr[i + 2] == 13 and ptr[i + 3] == 10 then
+      header_len = i + 2
+      break
     end
   end
+
+  if not header_len then
+    return nil
+  end
+
+  local content_length = get_content_length(ptr, 0, header_len)
+  if strbuffer.len(strbuf) < header_len + 2 + content_length then
+    return nil
+  end
+
+  strbuf:skip(header_len + 2) -- skip past header boundary
+  return strbuf:get(content_length)
 end
 
 --- @private
@@ -226,35 +235,25 @@ end
 function M.create_read_loop(handle_body, on_exit, on_error)
   on_exit = on_exit or function() end
   on_error = on_error or function() end
-  local co = coroutine.create(message_decoder)
-  coroutine.resume(co)
-  return function(err, chunk)
-    if err then
-      on_error(err, M.client_errors.READ_ERROR)
-      return
-    end
-
-    if not chunk then
-      on_exit()
-      return
-    end
-
-    if coroutine.status(co) == 'dead' then
-      return
-    end
-
-    while true do
-      local ok, res = coroutine.resume(co, chunk)
-      if not ok then
-        on_error(res, M.client_errors.INVALID_SERVER_MESSAGE)
-        break
-      elseif res then
-        handle_body(res)
-        chunk = ''
+  local message_stream = vim_transport.MessageStream.new(
+    message_decoder,
+    format_message_with_content_length,
+    function(err, chunk)
+      if err then
+        on_error(err, M.client_errors.READ_ERROR)
+      elseif chunk then
+        handle_body(chunk)
       else
-        break
+        on_exit()
       end
+    end,
+    function(err)
+      on_error(err, M.client_errors.INVALID_SERVER_MESSAGE)
     end
+  )
+
+  return function(err, chunk)
+    message_stream:feed(err, chunk)
   end
 end
 
@@ -263,7 +262,8 @@ end
 --- @field private message_index integer
 --- @field private message_callbacks table<integer, function> dict of message_id to callback
 --- @field private notify_reply_callbacks table<integer, function> dict of message_id to callback
---- @field private transport vim.lsp.rpc.Transport
+--- @field private transport vim.Transport
+--- @field private message_stream vim.MessageStream
 --- @field private dispatchers vim.lsp.rpc.Dispatchers
 ---
 --- See [vim.lsp.rpc.request()]
@@ -281,9 +281,11 @@ local Client = {}
 
 ---@package
 ---@param dispatchers vim.lsp.rpc.Dispatchers
----@param transport vim.lsp.rpc.Transport
+---@param transport vim.Transport
+---@param decode fun(buf: vim._core.stringbuffer): string?
+---@param format fun(msg: string): string
 ---@return vim.lsp.rpc.Client
-function Client.new(dispatchers, transport)
+function Client.new(dispatchers, transport, decode, format)
   local result = {
     message_index = 0,
     message_callbacks = {},
@@ -324,27 +326,26 @@ function Client.new(dispatchers, transport)
 
   ---@cast result vim.lsp.rpc.Client
   local self = setmetatable(result, { __index = Client })
-  --- @param body string
-  local function handle_body(body)
-    self:handle_body(body)
-  end
 
-  local function on_exit()
-    ---@diagnostic disable-next-line: invisible
-    self.transport:terminate()
-  end
-
-  --- @param errkind vim.lsp.rpc.ClientErrors
-  local function on_error(err, errkind)
-    self:on_error(errkind, err)
-    if errkind == M.client_errors.INVALID_SERVER_MESSAGE then
+  self.message_stream = vim_transport.MessageStream.new(decode, format, function(err, data)
+    if err then
+      self:on_error(M.client_errors.READ_ERROR, err)
+    elseif data then
+      self:handle_body(data)
+    else
       ---@diagnostic disable-next-line: invisible
       self.transport:terminate()
     end
-  end
+  end, function(err)
+    self:on_error(M.client_errors.INVALID_SERVER_MESSAGE, err)
+    ---@diagnostic disable-next-line: invisible
+    self.transport:terminate()
+  end)
 
-  local on_read = M.create_read_loop(handle_body, on_exit, on_error)
-  transport:listen(on_read, dispatchers.on_exit)
+  transport:listen(function(err, data)
+    ---@diagnostic disable-next-line: invisible
+    self.message_stream:feed(err, data)
+  end, dispatchers.on_exit)
   return self
 end
 
@@ -356,7 +357,7 @@ function Client:encode_and_send(payload)
   end
   local jsonstr = vim.json.encode(payload)
 
-  self.transport:write(format_message_with_content_length(jsonstr))
+  self.transport:write(self.message_stream.encode(jsonstr))
   return true
 end
 
@@ -636,13 +637,13 @@ function M.connect(host_or_path, port)
 
     dispatchers = merge_dispatchers(dispatchers)
 
-    local transport = lsp_transport.TransportConnect.new(host_or_path, port)
-    return Client.new(dispatchers, transport)
+    local transport = vim_transport.TransportConnect.new(host_or_path, port)
+    return Client.new(dispatchers, transport, message_decoder, format_message_with_content_length)
   end
 end
 
 --- Additional context for the LSP server process.
---- @class vim.lsp.rpc.ExtraSpawnParams
+--- @class vim.transport.ExtraSpawnParams
 --- @inlinedoc
 --- @field cwd? string Working directory for the LSP server process
 --- @field detached? boolean Detach the LSP server process from the current process
@@ -654,7 +655,7 @@ end
 ---
 --- @param cmd string[] Command to start the LSP server.
 --- @param dispatchers? vim.lsp.rpc.Dispatchers
---- @param extra_spawn_params? vim.lsp.rpc.ExtraSpawnParams
+--- @param extra_spawn_params? vim.transport.ExtraSpawnParams
 --- @return vim.lsp.rpc.Client
 function M.start(cmd, dispatchers, extra_spawn_params)
   log.info('Starting RPC client', { cmd = cmd, extra = extra_spawn_params })
@@ -664,8 +665,8 @@ function M.start(cmd, dispatchers, extra_spawn_params)
 
   dispatchers = merge_dispatchers(dispatchers)
 
-  local transport = lsp_transport.TransportRun.new(cmd, extra_spawn_params)
-  return Client.new(dispatchers, transport)
+  local transport = vim_transport.TransportRun.new(cmd, extra_spawn_params)
+  return Client.new(dispatchers, transport, message_decoder, format_message_with_content_length)
 end
 
 return M

runtime/lua/vim/net/_transport.lua

@@ -1,25 +1,24 @@
 local uv = vim.uv
 local log = require('vim.lsp.log')
+local strbuffer = require('vim._core.stringbuffer')
 
 --- Interface for transport implementations.
 ---
---- @class (private) vim.lsp.rpc.Transport
---- @field listen fun(self: vim.lsp.rpc.Transport, on_read: fun(err: any, data: string), on_exit: fun(code: integer, signal: integer))
---- @field write fun(self: vim.lsp.rpc.Transport, msg: string)
---- @field is_closing fun(self: vim.lsp.rpc.Transport): boolean
---- @field terminate fun(self: vim.lsp.rpc.Transport)
+--- @class (private, exact) vim.Transport
+--- @field listen fun(self: vim.Transport, on_read: fun(err: any, data: string), on_exit: fun(code: integer, signal: integer))
+--- @field write fun(self: vim.Transport, msg: string)
+--- @field is_closing fun(self: vim.Transport): boolean
+--- @field terminate fun(self: vim.Transport)
 
 --- Transport backed by newly spawned process using `vim.system()`.
 ---
---- @class (private) vim.lsp.rpc.Transport.Run : vim.lsp.rpc.Transport
---- @field cmd string[] Command to start the LSP server.
---- @field extra_spawn_params? vim.lsp.rpc.ExtraSpawnParams
---- @field sysobj? vim.SystemObj
+--- @class (private, exact) vim.TransportRun : vim.Transport
+--- @field private cmd string[] Command to start the process.
+--- @field private extra_spawn_params? vim.transport.ExtraSpawnParams
+--- @field private sysobj? vim.SystemObj
+--- @field new fun(cmd: string[], extra_spawn_params?: vim.transport.ExtraSpawnParams): vim.TransportRun
 local TransportRun = {}
 
---- @param cmd string[] Command to start the LSP server.
---- @param extra_spawn_params? vim.lsp.rpc.ExtraSpawnParams
---- @return vim.lsp.rpc.Transport.Run
 function TransportRun.new(cmd, extra_spawn_params)
   return setmetatable({
     cmd = cmd,
@@ -32,7 +31,7 @@ end
 function TransportRun:listen(on_read, on_exit)
   local function on_stderr(_, chunk)
     if chunk then
-      log.error('rpc', self.cmd[1], 'stderr', chunk)
+      log.error('transport', self.cmd[1], 'stderr', chunk)
     end
   end
 
@@ -64,10 +63,10 @@ function TransportRun:listen(on_read, on_exit)
   if not ok then ---@cast sysobj_or_err string
     local err = sysobj_or_err
     local sfx = err:match('ENOENT')
-        and '. The language server is either not installed, missing from PATH, or not executable.'
+        and '. The command is either not installed, missing from PATH, or not executable.'
       or string.format(' with error message: %s', err)
 
-    error(('Spawning language server with cmd: `%s` failed%s'):format(vim.inspect(self.cmd), sfx))
+    error(('Spawning process with cmd: `%s` failed%s'):format(vim.inspect(self.cmd), sfx))
   end ---@cast sysobj_or_err vim.SystemObj
 
   self.sysobj = sysobj_or_err
@@ -91,29 +90,27 @@ end
 
 --- Transport backed by an existing `uv.uv_pipe_t` or `uv.uv_tcp_t` connection.
 ---
---- @class (private) vim.lsp.rpc.Transport.Connect : vim.lsp.rpc.Transport
---- @field host_or_path string
---- @field port? integer
---- @field handle? uv.uv_pipe_t|uv.uv_tcp_t
+--- @class (private, exact) vim.TransportConnect : vim.Transport
+--- @field private host_or_path string
+--- @field private port? integer
+--- @field private handle? uv.uv_pipe_t|uv.uv_tcp_t
 --- Connect returns a PublicClient synchronously so the caller
 --- can immediately send messages before the connection is established.
 --- These messages are buffered in `msgbuf`.
---- @field connected boolean
---- @field closing boolean
---- @field msgbuf vim.Ringbuf
---- @field on_exit? fun(code: integer, signal: integer)
+--- @field private connected boolean
+--- @field private closing boolean
+--- @field private msgbuf vim.Ringbuf
+--- @field private on_exit? fun(code: integer, signal: integer)
+--- @field new fun(host_or_path: string, port?: integer): vim.TransportConnect
 local TransportConnect = {}
 
---- @param host_or_path string
---- @param port? integer
---- @return vim.lsp.rpc.Transport.Connect
 function TransportConnect.new(host_or_path, port)
   return setmetatable({
     host_or_path = host_or_path,
     port = port,
     connected = false,
     -- size should be enough because the client can't really do anything until initialization is done
-    -- which required a response from the server - implying the connection got established
+    -- which required a response from the process - implying the connection got established
     msgbuf = vim.ringbuf(10),
     closing = false,
   }, { __index = TransportConnect })
@@ -186,7 +183,65 @@ function TransportConnect:terminate()
   end
 end
 
+--- Create a message stream from a decoder.
+---
+--- The decoder consumes from the given string buffer
+--- and returns a message body when a full message is available.
+--- `nil` means it needs more transport data.
+--- decoder errors are reported through `on_error`.
+---
+---@class (private, exact) vim.MessageStream
+---@field private strbuf string.buffer
+---@field private decode fun(strbuf: string.buffer): string?
+---@field private on_read fun(err: string?, data: string?)
+---@field private on_error fun(err: any)
+---@field feed fun(self: vim.MessageStream, err: string?, data: string?)
+---@field encode fun(msg: string): string
+---@field new fun(decode: (fun(strbuf: string.buffer): string?), encode: (fun(msg: string): string), on_read: fun(err: string?, data: string?), on_error: fun(err: any)): vim.MessageStream
+local MessageStream = {}
+
+---@param decode fun(strbuf: string.buffer): string?
+---@param encode fun(msg: string): string
+---@param on_read fun(err: string?, data: string?)
+---@param on_error fun(err: any)
+---@return vim.MessageStream
+function MessageStream.new(decode, encode, on_read, on_error)
+  return setmetatable({
+    strbuf = strbuffer.new(),
+    decode = decode,
+    on_read = on_read,
+    on_error = on_error,
+    encode = encode,
+  }, { __index = MessageStream })
+end
+
+---@param err string?
+---@param data string?
+function MessageStream:feed(err, data)
+  if err then
+    self.on_read(err, nil)
+    return
+  elseif data == nil then
+    self.on_read(nil, nil)
+    return
+  end
+
+  self.strbuf:put(data)
+
+  while true do
+    local ok, body = pcall(self.decode, self.strbuf)
+    if not ok then
+      self.on_error(body)
+      return
+    elseif body == nil then
+      break
+    end
+    self.on_read(nil, body)
+  end
+end
+
 return {
   TransportRun = TransportRun,
   TransportConnect = TransportConnect,
+  MessageStream = MessageStream,
 }