refactor(lsp): simplify client tracking

- Remove:
    - uninitialized_clients
    - active_clients
    - all_buffer_active_clients
- Add:
    - all_clients

- Use `lsp.get_clients()` to get buffer clients.

runtime/doc/lsp.txt

@@ -907,8 +907,8 @@ stop_client({client_id}, {force})                      *vim.lsp.stop_client()*
     for this client, then force-shutdown is attempted.
 
     Parameters: ~
-      • {client_id}  (`integer|vim.lsp.Client`) id or |vim.lsp.Client| object,
-                     or list thereof
+      • {client_id}  (`integer|integer[]|vim.lsp.Client[]`) id, list of id's,
+                     or list of |vim.lsp.Client| objects
       • {force}      (`boolean?`) shutdown forcefully
 
 tagfunc({pattern}, {flags})                                *vim.lsp.tagfunc()*

runtime/lua/vim/lsp.lua

@@ -1,7 +1,5 @@
 local api = vim.api
-local tbl_isempty, tbl_extend = vim.tbl_isempty, vim.tbl_extend
 local validate = vim.validate
-local if_nil = vim.F.if_nil
 
 local lsp = vim._defer_require('vim.lsp', {
   _changetracking = ..., --- @module 'vim.lsp._changetracking'
@@ -108,40 +106,7 @@ function lsp._buf_get_line_ending(bufnr)
 end
 
 -- Tracks all clients created via lsp.start_client
-local active_clients = {} --- @type table<integer,vim.lsp.Client>
-local all_buffer_active_clients = {} --- @type table<integer,table<integer,true>>
-local uninitialized_clients = {} --- @type table<integer,vim.lsp.Client>
-
----@param bufnr? integer
----@param fn fun(client: vim.lsp.Client, client_id: integer, bufnr: integer)
-local function for_each_buffer_client(bufnr, fn, restrict_client_ids)
-  validate({
-    fn = { fn, 'f' },
-    restrict_client_ids = { restrict_client_ids, 't', true },
-  })
-  bufnr = resolve_bufnr(bufnr)
-  local client_ids = all_buffer_active_clients[bufnr]
-  if not client_ids or tbl_isempty(client_ids) then
-    return
-  end
-
-  if restrict_client_ids and #restrict_client_ids > 0 then
-    local filtered_client_ids = {} --- @type table<integer,true>
-    for client_id in pairs(client_ids) do
-      if vim.list_contains(restrict_client_ids, client_id) then
-        filtered_client_ids[client_id] = true
-      end
-    end
-    client_ids = filtered_client_ids
-  end
-
-  for client_id in pairs(client_ids) do
-    local client = active_clients[client_id]
-    if client then
-      fn(client, client_id, bufnr)
-    end
-  end
-end
+local all_clients = {} --- @type table<integer,vim.lsp.Client>
 
 local client_errors_base = table.maxn(lsp.rpc.client_errors)
 local client_errors_offset = 0
@@ -156,7 +121,7 @@ end
 --- Can be used to look up the string from a the number or the number
 --- from the string.
 --- @nodoc
-lsp.client_errors = tbl_extend(
+lsp.client_errors = vim.tbl_extend(
   'error',
   lsp.rpc.client_errors,
   client_error('BEFORE_INIT_CALLBACK_ERROR'),
@@ -258,14 +223,12 @@ function lsp.start(config, opts)
 
   local bufnr = resolve_bufnr(opts.bufnr)
 
-  for _, clients in ipairs({ uninitialized_clients, lsp.get_clients() }) do
-    for _, client in pairs(clients) do
+  for _, client in pairs(all_clients) do
     if reuse_client(client, config) then
       lsp.buf_attach_client(bufnr, client.id)
       return client.id
     end
   end
-  end
 
   local client_id = lsp.start_client(config)
 
@@ -383,30 +346,13 @@ local function reset_defaults(bufnr)
   end)
 end
 
---- @param client vim.lsp.Client
-local function on_client_init(client)
-  local id = client.id
-  uninitialized_clients[id] = nil
-  -- Only assign after initialized.
-  active_clients[id] = client
-  -- If we had been registered before we start, then send didOpen This can
-  -- happen if we attach to buffers before initialize finishes or if
-  -- someone restarts a client.
-  for bufnr, client_ids in pairs(all_buffer_active_clients) do
-    if client_ids[id] then
-      client.on_attach(bufnr)
-    end
-  end
-end
-
 --- @param code integer
 --- @param signal integer
 --- @param client_id integer
 local function on_client_exit(code, signal, client_id)
-  local client = active_clients[client_id] or uninitialized_clients[client_id]
+  local client = all_clients[client_id]
 
-  for bufnr, client_ids in pairs(all_buffer_active_clients) do
-    if client_ids[client_id] then
+  for bufnr in pairs(client.attached_buffers) do
     vim.schedule(function()
       if client and client.attached_buffers[bufnr] then
         api.nvim_exec_autocmds('LspDetach', {
@@ -418,22 +364,20 @@ local function on_client_exit(code, signal, client_id)
 
       local namespace = vim.lsp.diagnostic.get_namespace(client_id)
       vim.diagnostic.reset(namespace, bufnr)
+      client.attached_buffers[bufnr] = nil
 
-        client_ids[client_id] = nil
-        if vim.tbl_isempty(client_ids) then
+      if #lsp.get_clients({ bufnr = bufnr, _uninitialized = true }) == 0 then
         reset_defaults(bufnr)
       end
     end)
   end
-  end
 
   local name = client.name or 'unknown'
 
   -- Schedule the deletion of the client object so that it exists in the execution of LspDetach
   -- autocommands
   vim.schedule(function()
-    active_clients[client_id] = nil
-    uninitialized_clients[client_id] = nil
+    all_clients[client_id] = nil
 
     -- Client can be absent if executable starts, but initialize fails
     -- init/attach won't have happened
@@ -465,13 +409,10 @@ function lsp.start_client(config)
     return
   end
 
-  --- @diagnostic disable-next-line: invisible
-  table.insert(client._on_init_cbs, on_client_init)
   --- @diagnostic disable-next-line: invisible
   table.insert(client._on_exit_cbs, on_client_exit)
 
-  -- Store the uninitialized_clients for cleanup in case we exit before initialize finishes.
-  uninitialized_clients[client.id] = client
+  all_clients[client.id] = client
 
   client:initialize()
 
@@ -495,7 +436,7 @@ local function text_document_did_change_handler(
   new_lastline
 )
   -- Detach (nvim_buf_attach) via returning True to on_lines if no clients are attached
-  if tbl_isempty(all_buffer_active_clients[bufnr] or {}) then
+  if #lsp.get_clients({ bufnr = bufnr }) == 0 then
     return true
   end
   util.buf_versions[bufnr] = changedtick
@@ -543,30 +484,9 @@ local function text_document_did_save_handler(bufnr)
   end
 end
 
---- Implements the `textDocument/did…` notifications required to track a buffer
---- for any language server.
----
---- Without calling this, the server won't be notified of changes to a buffer.
----
----@param bufnr (integer) Buffer handle, or 0 for current
----@param client_id (integer) Client id
----@return boolean success `true` if client was attached successfully; `false` otherwise
-function lsp.buf_attach_client(bufnr, client_id)
-  validate({
-    bufnr = { bufnr, 'n', true },
-    client_id = { client_id, 'n' },
-  })
-  bufnr = resolve_bufnr(bufnr)
-  if not api.nvim_buf_is_loaded(bufnr) then
-    log.warn(string.format('buf_attach_client called on unloaded buffer (id: %d): ', bufnr))
-    return false
-  end
-  local buffer_client_ids = all_buffer_active_clients[bufnr]
-  -- This is our first time attaching to this buffer.
-  if not buffer_client_ids then
-    buffer_client_ids = {}
-    all_buffer_active_clients[bufnr] = buffer_client_ids
-
+--- @param bufnr integer
+--- @param client_id integer
+local function buf_attach(bufnr, client_id)
   local uri = vim.uri_from_bufnr(bufnr)
   local augroup = ('lsp_c_%d_b_%d_save'):format(client_id, bufnr)
   local group = api.nvim_create_augroup(augroup, { clear = true })
@@ -625,10 +545,11 @@ function lsp.buf_attach_client(bufnr, client_id)
         if vim.tbl_get(client.server_capabilities, 'textDocumentSync', 'openClose') then
           client.notify(ms.textDocument_didClose, params)
         end
+      end
+      for _, client in ipairs(all_clients) do
         client.attached_buffers[bufnr] = nil
       end
       util.buf_versions[bufnr] = nil
-        all_buffer_active_clients[bufnr] = nil
     end,
     -- TODO if we know all of the potential clients ahead of time, then we
     -- could conditionally set this.
@@ -637,16 +558,44 @@ function lsp.buf_attach_client(bufnr, client_id)
   })
 end
 
-  if buffer_client_ids[client_id] then
+--- Implements the `textDocument/did…` notifications required to track a buffer
+--- for any language server.
+---
+--- Without calling this, the server won't be notified of changes to a buffer.
+---
+---@param bufnr (integer) Buffer handle, or 0 for current
+---@param client_id (integer) Client id
+---@return boolean success `true` if client was attached successfully; `false` otherwise
+function lsp.buf_attach_client(bufnr, client_id)
+  validate({
+    bufnr = { bufnr, 'n', true },
+    client_id = { client_id, 'n' },
+  })
+  bufnr = resolve_bufnr(bufnr)
+  if not api.nvim_buf_is_loaded(bufnr) then
+    log.warn(string.format('buf_attach_client called on unloaded buffer (id: %d): ', bufnr))
+    return false
+  end
+  -- This is our first time attaching to this buffer.
+  if #lsp.get_clients({ bufnr = bufnr, _uninitialized = true }) == 0 then
+    buf_attach(bufnr, client_id)
+  end
+
+  local client = lsp.get_client_by_id(client_id)
+  if not client then
+    return false
+  end
+
+  if client.attached_buffers[bufnr] then
     return true
   end
-  -- This is our first time attaching this client to this buffer.
-  buffer_client_ids[client_id] = true
 
-  local client = active_clients[client_id]
+  client.attached_buffers[bufnr] = true
+
+  -- This is our first time attaching this client to this buffer.
   -- Send didOpen for the client if it is initialized. If it isn't initialized
   -- then it will send didOpen on initialize.
-  if client then
+  if client.initialized then
     client:_on_attach(bufnr)
   end
   return true
@@ -665,7 +614,7 @@ function lsp.buf_detach_client(bufnr, client_id)
   })
   bufnr = resolve_bufnr(bufnr)
 
-  local client = lsp.get_client_by_id(client_id)
+  local client = all_clients[client_id]
   if not client or not client.attached_buffers[bufnr] then
     vim.notify(
       string.format(
@@ -694,11 +643,6 @@ function lsp.buf_detach_client(bufnr, client_id)
   client.attached_buffers[bufnr] = nil
   util.buf_versions[bufnr] = nil
 
-  all_buffer_active_clients[bufnr][client_id] = nil
-  if #vim.tbl_keys(all_buffer_active_clients[bufnr]) == 0 then
-    all_buffer_active_clients[bufnr] = nil
-  end
-
   local namespace = lsp.diagnostic.get_namespace(client_id)
   vim.diagnostic.reset(namespace, bufnr)
 end
@@ -708,7 +652,7 @@ end
 ---@param bufnr (integer) Buffer handle, or 0 for current
 ---@param client_id (integer) the client id
 function lsp.buf_is_attached(bufnr, client_id)
-  return (all_buffer_active_clients[resolve_bufnr(bufnr)] or {})[client_id] == true
+  return lsp.get_clients({ bufnr = bufnr, id = client_id, _uninitialized = true })[1] ~= nil
 end
 
 --- Gets a client by id, or nil if the id is invalid.
@@ -718,7 +662,7 @@ end
 ---
 ---@return (nil|vim.lsp.Client) client rpc object
 function lsp.get_client_by_id(client_id)
-  return active_clients[client_id] or uninitialized_clients[client_id]
+  return all_clients[client_id]
 end
 
 --- Returns list of buffers attached to client_id.
@@ -726,7 +670,7 @@ end
 ---@param client_id integer client id
 ---@return integer[] buffers list of buffer ids
 function lsp.get_buffers_by_client_id(client_id)
-  local client = lsp.get_client_by_id(client_id)
+  local client = all_clients[client_id]
   return client and vim.tbl_keys(client.attached_buffers) or {}
 end
 
@@ -742,17 +686,22 @@ end
 --- By default asks the server to shutdown, unless stop was requested
 --- already for this client, then force-shutdown is attempted.
 ---
----@param client_id integer|vim.lsp.Client id or |vim.lsp.Client| object, or list thereof
----@param force boolean|nil shutdown forcefully
+---@param client_id integer|integer[]|vim.lsp.Client[] id, list of id's, or list of |vim.lsp.Client| objects
+---@param force? boolean shutdown forcefully
 function lsp.stop_client(client_id, force)
+  --- @type integer[]|vim.lsp.Client[]
   local ids = type(client_id) == 'table' and client_id or { client_id }
   for _, id in ipairs(ids) do
-    if type(id) == 'table' and id.stop ~= nil then
+    if type(id) == 'table' then
+      if id.stop then
         id.stop(force)
-    elseif active_clients[id] then
-      active_clients[id].stop(force)
-    elseif uninitialized_clients[id] then
-      uninitialized_clients[id].stop(true)
+      end
+    else
+      --- @cast id -vim.lsp.Client
+      local client = all_clients[id]
+      if client then
+        client.stop(force)
+      end
     end
   end
 end
@@ -772,6 +721,9 @@ end
 ---
 --- Only return clients supporting the given method
 --- @field method? string
+---
+--- Also return uninitialized clients.
+--- @field package _uninitialized? boolean
 
 --- Get active clients.
 ---
@@ -784,15 +736,16 @@ function lsp.get_clients(filter)
 
   local clients = {} --- @type vim.lsp.Client[]
 
-  local t = filter.bufnr and (all_buffer_active_clients[resolve_bufnr(filter.bufnr)] or {})
-    or active_clients
-  for client_id in pairs(t) do
-    local client = active_clients[client_id]
+  local bufnr = filter.bufnr and resolve_bufnr(filter.bufnr)
+
+  for _, client in pairs(all_clients) do
     if
       client
       and (filter.id == nil or client.id == filter.id)
+      and (filter.bufnr == nil or client.attached_buffers[bufnr])
       and (filter.name == nil or client.name == filter.name)
       and (filter.method == nil or client.supports_method(filter.method, { bufnr = filter.bufnr }))
+      and (filter._uninitialized or client.initialized)
     then
       clients[#clients + 1] = client
     end
@@ -810,15 +763,9 @@ end
 api.nvim_create_autocmd('VimLeavePre', {
   desc = 'vim.lsp: exit handler',
   callback = function()
+    local active_clients = lsp.get_clients()
     log.info('exit_handler', active_clients)
-    for _, client in pairs(uninitialized_clients) do
-      client.stop(true)
-    end
-    -- TODO handle v:dying differently?
-    if tbl_isempty(active_clients) then
-      return
-    end
-    for _, client in pairs(active_clients) do
+    for _, client in pairs(all_clients) do
       client.stop()
     end
 
@@ -827,7 +774,7 @@ api.nvim_create_autocmd('VimLeavePre', {
     local send_kill = false
 
     for client_id, client in pairs(active_clients) do
-      local timeout = if_nil(client.flags.exit_timeout, false)
+      local timeout = client.flags.exit_timeout
       if timeout then
         send_kill = true
         timeouts[client_id] = timeout
@@ -910,7 +857,7 @@ function lsp.buf_request(bufnr, method, params, handler)
 
   local function _cancel_all_requests()
     for client_id, request_id in pairs(client_request_ids) do
-      local client = active_clients[client_id]
+      local client = all_clients[client_id]
       client.cancel_request(request_id)
     end
   end
@@ -1106,7 +1053,7 @@ end
 ---@return boolean stopped true if client is stopped, false otherwise.
 function lsp.client_is_stopped(client_id)
   assert(client_id, 'missing client_id param')
-  return active_clients[client_id] == nil and not uninitialized_clients[client_id]
+  return not all_clients[client_id]
 end
 
 --- Gets a map of client_id:client pairs for the given buffer, where each value
@@ -1172,7 +1119,11 @@ function lsp.for_each_buffer_client(bufnr, fn)
     'lsp.get_clients({ bufnr = bufnr }) with regular loop',
     '0.12'
   )
-  return for_each_buffer_client(bufnr, fn)
+  bufnr = resolve_bufnr(bufnr)
+
+  for _, client in pairs(lsp.get_clients({ bufnr = bufnr })) do
+    fn(client, client.id, bufnr)
+  end
 end
 
 --- Function to manage overriding defaults for LSP handlers.

runtime/lua/vim/lsp/client.lua

@@ -185,6 +185,10 @@ local validate = vim.validate
 --- @field root_dir string
 ---
 --- @field attached_buffers table<integer,true>
+---
+--- Buffers that should be attached to upon initialize()
+--- @field package _buffers_to_attach table<integer,true>
+---
 --- @field private _log_prefix string
 ---
 --- Track this so that we can escalate automatically if we've already tried a
@@ -608,8 +612,16 @@ function Client:initialize()
       self:_notify(ms.workspace_didChangeConfiguration, { settings = self.settings })
     end
 
+    -- If server is being restarted, make sure to re-attach to any previously attached buffers.
+    -- Save which buffers before on_init in case new buffers are attached.
+    local reattach_bufs = vim.deepcopy(self.attached_buffers)
+
     self:_run_callbacks(self._on_init_cbs, lsp.client_errors.ON_INIT_CALLBACK_ERROR, self, result)
 
+    for buf in pairs(reattach_bufs) do
+      self:_on_attach(buf)
+    end
+
     log.info(
       self._log_prefix,
       'server_capabilities',
@@ -761,7 +773,7 @@ function Client:_request_sync(method, params, timeout_ms, bufnr)
   return request_result
 end
 
---- @private
+--- @package
 --- Sends a notification to an LSP server.
 ---
 --- @param method string LSP method name.

test/functional/fixtures/fake-lsp-server.lua

@@ -983,7 +983,7 @@ local test_name = arg[1]
 local timeout = arg[2]
 assert(type(test_name) == 'string', 'test_name must be specified as first arg.')
 
-local kill_timer = vim.uv.new_timer()
+local kill_timer = assert(vim.uv.new_timer())
 kill_timer:start(timeout or 1e3, 0, function()
   kill_timer:stop()
   kill_timer:close()

test/functional/plugin/lsp_spec.lua

@@ -282,8 +282,7 @@ describe('LSP', function()
       local client --- @type vim.lsp.Client
       test_rpc_server {
         test_name = 'basic_finish',
-        on_init = function(_client)
-          client = _client
+        on_setup = function()
           exec_lua [[
             BUFFER = vim.api.nvim_create_buf(false, true)
           ]]
@@ -292,6 +291,9 @@ describe('LSP', function()
           exec_lua [[
             vim.api.nvim_command(BUFFER.."bwipeout")
           ]]
+        end,
+        on_init = function(_client)
+          client = _client
           client.notify('finish')
         end,
         on_exit = function(code, signal)
@@ -806,14 +808,12 @@ describe('LSP', function()
               BUFFER = vim.api.nvim_get_current_buf()
               lsp.buf_attach_client(BUFFER, TEST_RPC_CLIENT_ID)
               vim.lsp.handlers['textDocument/typeDefinition'] = function() end
+              vim.cmd(BUFFER.."bwipeout")
             ]=])
         end,
         on_init = function(client)
           client.stop()
           exec_lua('vim.lsp.buf.type_definition()')
-          exec_lua [[
-            vim.api.nvim_command(BUFFER.."bwipeout")
-          ]]
         end,
         on_exit = function(code, signal)
           eq(0, code, 'exit code')
@@ -1058,21 +1058,21 @@ describe('LSP', function()
       local client --- @type vim.lsp.Client
       test_rpc_server {
         test_name = 'basic_finish',
-        on_init = function(_client)
-          client = _client
+        on_setup = function()
           exec_lua [[
             BUFFER = vim.api.nvim_create_buf(false, true)
             vim.api.nvim_buf_set_lines(BUFFER, 0, -1, false, {
               "testing";
               "123";
             })
+            assert(TEST_RPC_CLIENT_ID == 1)
+            assert(lsp.buf_attach_client(BUFFER, TEST_RPC_CLIENT_ID))
+            assert(lsp.buf_is_attached(BUFFER, TEST_RPC_CLIENT_ID))
+            vim.cmd(BUFFER.."bwipeout")
           ]]
-          eq(1, exec_lua('return TEST_RPC_CLIENT_ID'))
-          eq(true, exec_lua('return lsp.buf_attach_client(BUFFER, TEST_RPC_CLIENT_ID)'))
-          eq(true, exec_lua('return lsp.buf_is_attached(BUFFER, TEST_RPC_CLIENT_ID)'))
-          exec_lua [[
-            vim.api.nvim_command(BUFFER.."bwipeout")
-          ]]
+        end,
+        on_init = function(_client)
+          client = _client
           local full_kind = exec_lua("return require'vim.lsp.protocol'.TextDocumentSyncKind.Full")
           eq(full_kind, client.server_capabilities().textDocumentSync.change)
           eq(true, client.server_capabilities().textDocumentSync.openClose)