From 69806560f2d1a9f12915d34700c329eadf843170 Mon Sep 17 00:00:00 2001
From: Vinicius Stock <vinicius.stock@shopify.com>
Date: Wed, 11 Dec 2024 17:06:29 -0500
Subject: [PATCH] Locate targets under mutex lock

---
 lib/ruby_lsp/base_server.rb                   | 13 +++--
 lib/ruby_lsp/global_state.rb                  |  6 +++
 lib/ruby_lsp/requests/code_action_resolve.rb  | 48 +++++++++++--------
 lib/ruby_lsp/requests/completion.rb           |  4 +-
 lib/ruby_lsp/requests/definition.rb           |  4 +-
 lib/ruby_lsp/requests/document_highlight.rb   |  4 +-
 lib/ruby_lsp/requests/hover.rb                |  4 +-
 lib/ruby_lsp/requests/prepare_rename.rb       |  8 +++-
 lib/ruby_lsp/requests/references.rb           |  4 +-
 lib/ruby_lsp/requests/rename.rb               |  4 +-
 lib/ruby_lsp/requests/show_syntax_tree.rb     | 14 ++++--
 lib/ruby_lsp/requests/signature_help.rb       |  5 +-
 lib/ruby_lsp/server.rb                        | 11 +++--
 .../prepare_rename_expectations_test.rb       |  3 +-
 14 files changed, 86 insertions(+), 46 deletions(-)

diff --git a/lib/ruby_lsp/base_server.rb b/lib/ruby_lsp/base_server.rb
index 8fac3592c..03a94b16b 100644
--- a/lib/ruby_lsp/base_server.rb
+++ b/lib/ruby_lsp/base_server.rb
@@ -18,22 +18,21 @@ def initialize(**options)
       @incoming_queue = T.let(Thread::Queue.new, Thread::Queue)
       @outgoing_queue = T.let(Thread::Queue.new, Thread::Queue)
       @cancelled_requests = T.let([], T::Array[Integer])
-      @mutex = T.let(Mutex.new, Mutex)
       @worker = T.let(new_worker, Thread)
       @current_request_id = T.let(1, Integer)
       @store = T.let(Store.new, Store)
+      @global_state = T.let(GlobalState.new, GlobalState)
       @outgoing_dispatcher = T.let(
         Thread.new do
           unless @test_mode
             while (message = @outgoing_queue.pop)
-              @mutex.synchronize { @writer.write(message.to_hash) }
+              @global_state.synchronize { @writer.write(message.to_hash) }
             end
           end
         end,
         Thread,
       )
 
-      @global_state = T.let(GlobalState.new, GlobalState)
       Thread.main.priority = 1
 
       # We read the initialize request in `exe/ruby-lsp` to be able to determine the workspace URI where Bundler should
@@ -51,7 +50,7 @@ def start
         # source. Altering the source reference during parsing will put the parser in an invalid internal state, since
         # it started parsing with one source but then it changed in the middle. We don't want to do this for text
         # synchronization notifications
-        @mutex.synchronize do
+        @global_state.synchronize do
           uri = message.dig(:params, :textDocument, :uri)
 
           if uri
@@ -95,14 +94,14 @@ def start
              "$/cancelRequest"
           process_message(message)
         when "shutdown"
-          @mutex.synchronize do
+          @global_state.synchronize do
             send_log_message("Shutting down Ruby LSP...")
             shutdown
             run_shutdown
             @writer.write(Result.new(id: message[:id], response: nil).to_hash)
           end
         when "exit"
-          @mutex.synchronize do
+          @global_state.synchronize do
             status = @incoming_queue.closed? ? 0 : 1
             send_log_message("Shutdown complete with status #{status}")
             exit(status)
@@ -157,7 +156,7 @@ def new_worker
           id = message[:id]
 
           # Check if the request was cancelled before trying to process it
-          @mutex.synchronize do
+          @global_state.synchronize do
             if id && @cancelled_requests.include?(id)
               send_message(Error.new(
                 id: id,
diff --git a/lib/ruby_lsp/global_state.rb b/lib/ruby_lsp/global_state.rb
index 004df40d9..2c6687eb0 100644
--- a/lib/ruby_lsp/global_state.rb
+++ b/lib/ruby_lsp/global_state.rb
@@ -53,6 +53,12 @@ def initialize
       )
       @client_capabilities = T.let(ClientCapabilities.new, ClientCapabilities)
       @enabled_feature_flags = T.let({}, T::Hash[Symbol, T::Boolean])
+      @mutex = T.let(Mutex.new, Mutex)
+    end
+
+    sig { type_parameters(:T).params(block: T.proc.returns(T.type_parameter(:T))).returns(T.type_parameter(:T)) }
+    def synchronize(&block)
+      @mutex.synchronize(&block)
     end
 
     sig { params(addon_name: String).returns(T.nilable(T::Hash[Symbol, T.untyped])) }
diff --git a/lib/ruby_lsp/requests/code_action_resolve.rb b/lib/ruby_lsp/requests/code_action_resolve.rb
index 78b545db1..4c467e384 100644
--- a/lib/ruby_lsp/requests/code_action_resolve.rb
+++ b/lib/ruby_lsp/requests/code_action_resolve.rb
@@ -54,10 +54,12 @@ def switch_block_style
         source_range = @code_action.dig(:data, :range)
         return Error::EmptySelection if source_range[:start] == source_range[:end]
 
-        target = @document.locate_first_within_range(
-          @code_action.dig(:data, :range),
-          node_types: [Prism::CallNode],
-        )
+        target = @global_state.synchronize do
+          @document.locate_first_within_range(
+            @code_action.dig(:data, :range),
+            node_types: [Prism::CallNode],
+          )
+        end
 
         return Error::InvalidTargetRange unless target.is_a?(Prism::CallNode)
 
@@ -92,10 +94,13 @@ def refactor_variable
         source_range = @code_action.dig(:data, :range)
         return Error::EmptySelection if source_range[:start] == source_range[:end]
 
-        scanner = @document.create_scanner
-        start_index = scanner.find_char_position(source_range[:start])
-        end_index = scanner.find_char_position(source_range[:end])
-        extracted_source = T.must(@document.source[start_index...end_index])
+        start_index, extracted_source = @global_state.synchronize do
+          scanner = @document.create_scanner
+          start_index = scanner.find_char_position(source_range[:start])
+          end_index = scanner.find_char_position(source_range[:end])
+          extracted_source = T.must(@document.source[start_index...end_index])
+          [start_index, extracted_source]
+        end
 
         # Find the closest statements node, so that we place the refactor in a valid position
         node_context = RubyDocument
@@ -192,10 +197,13 @@ def refactor_method
         source_range = @code_action.dig(:data, :range)
         return Error::EmptySelection if source_range[:start] == source_range[:end]
 
-        scanner = @document.create_scanner
-        start_index = scanner.find_char_position(source_range[:start])
-        end_index = scanner.find_char_position(source_range[:end])
-        extracted_source = T.must(@document.source[start_index...end_index])
+        start_index, extracted_source = @global_state.synchronize do
+          scanner = @document.create_scanner
+          start_index = scanner.find_char_position(source_range[:start])
+          end_index = scanner.find_char_position(source_range[:end])
+          extracted_source = T.must(@document.source[start_index...end_index])
+          [start_index, extracted_source]
+        end
 
         # Find the closest method declaration node, so that we place the refactor in a valid position
         node_context = RubyDocument.locate(
@@ -305,13 +313,15 @@ def recursively_switch_nested_block_styles(node, indentation)
       def switch_block_body(body, indentation)
         # Check if there are any nested blocks inside of the current block
         body_loc = body.location
-        nested_block = @document.locate_first_within_range(
-          {
-            start: { line: body_loc.start_line - 1, character: body_loc.start_column },
-            end: { line: body_loc.end_line - 1, character: body_loc.end_column },
-          },
-          node_types: [Prism::BlockNode],
-        )
+        nested_block = @global_state.synchronize do
+          @document.locate_first_within_range(
+            {
+              start: { line: body_loc.start_line - 1, character: body_loc.start_column },
+              end: { line: body_loc.end_line - 1, character: body_loc.end_column },
+            },
+            node_types: [Prism::BlockNode],
+          )
+        end
 
         body_content = body.slice.dup
 
diff --git a/lib/ruby_lsp/requests/completion.rb b/lib/ruby_lsp/requests/completion.rb
index 5ed77c0df..8f3f3c82f 100644
--- a/lib/ruby_lsp/requests/completion.rb
+++ b/lib/ruby_lsp/requests/completion.rb
@@ -40,7 +40,9 @@ def initialize(document, global_state, params, sorbet_level, dispatcher)
         @dispatcher = dispatcher
         # Completion always receives the position immediately after the character that was just typed. Here we adjust it
         # back by 1, so that we find the right node
-        char_position = document.create_scanner.find_char_position(params[:position]) - 1
+        char_position = global_state.synchronize do
+          document.create_scanner.find_char_position(params[:position]) - 1
+        end
         delegate_request_if_needed!(global_state, document, char_position)
 
         node_context = RubyDocument.locate(
diff --git a/lib/ruby_lsp/requests/definition.rb b/lib/ruby_lsp/requests/definition.rb
index e4817e3db..a10d40b07 100644
--- a/lib/ruby_lsp/requests/definition.rb
+++ b/lib/ruby_lsp/requests/definition.rb
@@ -29,7 +29,9 @@ def initialize(document, global_state, position, dispatcher, sorbet_level)
         )
         @dispatcher = dispatcher
 
-        char_position = document.create_scanner.find_char_position(position)
+        char_position = global_state.synchronize do
+          document.create_scanner.find_char_position(position)
+        end
         delegate_request_if_needed!(global_state, document, char_position)
 
         node_context = RubyDocument.locate(
diff --git a/lib/ruby_lsp/requests/document_highlight.rb b/lib/ruby_lsp/requests/document_highlight.rb
index e2d7dcb33..9e9af9fc4 100644
--- a/lib/ruby_lsp/requests/document_highlight.rb
+++ b/lib/ruby_lsp/requests/document_highlight.rb
@@ -25,7 +25,9 @@ class DocumentHighlight < Request
       end
       def initialize(global_state, document, position, dispatcher)
         super()
-        char_position = document.create_scanner.find_char_position(position)
+        char_position = global_state.synchronize do
+          document.create_scanner.find_char_position(position)
+        end
         delegate_request_if_needed!(global_state, document, char_position)
 
         node_context = RubyDocument.locate(
diff --git a/lib/ruby_lsp/requests/hover.rb b/lib/ruby_lsp/requests/hover.rb
index 55c5ea8ba..908e349a0 100644
--- a/lib/ruby_lsp/requests/hover.rb
+++ b/lib/ruby_lsp/requests/hover.rb
@@ -34,7 +34,9 @@ def provider
       def initialize(document, global_state, position, dispatcher, sorbet_level)
         super()
 
-        char_position = document.create_scanner.find_char_position(position)
+        char_position = global_state.synchronize do
+          document.create_scanner.find_char_position(position)
+        end
         delegate_request_if_needed!(global_state, document, char_position)
 
         node_context = RubyDocument.locate(
diff --git a/lib/ruby_lsp/requests/prepare_rename.rb b/lib/ruby_lsp/requests/prepare_rename.rb
index d49c60f03..90d2a0a75 100644
--- a/lib/ruby_lsp/requests/prepare_rename.rb
+++ b/lib/ruby_lsp/requests/prepare_rename.rb
@@ -12,19 +12,23 @@ class PrepareRename < Request
 
       sig do
         params(
+          global_state: GlobalState,
           document: RubyDocument,
           position: T::Hash[Symbol, T.untyped],
         ).void
       end
-      def initialize(document, position)
+      def initialize(global_state, document, position)
         super()
+        @global_state = global_state
         @document = document
         @position = T.let(position, T::Hash[Symbol, Integer])
       end
 
       sig { override.returns(T.nilable(Interface::Range)) }
       def perform
-        char_position = @document.create_scanner.find_char_position(@position)
+        char_position = @global_state.synchronize do
+          @document.create_scanner.find_char_position(@position)
+        end
 
         node_context = RubyDocument.locate(
           @document.parse_result.value,
diff --git a/lib/ruby_lsp/requests/references.rb b/lib/ruby_lsp/requests/references.rb
index bd7694f37..11b10b2d7 100644
--- a/lib/ruby_lsp/requests/references.rb
+++ b/lib/ruby_lsp/requests/references.rb
@@ -30,7 +30,9 @@ def initialize(global_state, store, document, params)
       sig { override.returns(T::Array[Interface::Location]) }
       def perform
         position = @params[:position]
-        char_position = @document.create_scanner.find_char_position(position)
+        char_position = @global_state.synchronize do
+          @document.create_scanner.find_char_position(position)
+        end
 
         node_context = RubyDocument.locate(
           @document.parse_result.value,
diff --git a/lib/ruby_lsp/requests/rename.rb b/lib/ruby_lsp/requests/rename.rb
index 0825476f8..ef1fea26f 100644
--- a/lib/ruby_lsp/requests/rename.rb
+++ b/lib/ruby_lsp/requests/rename.rb
@@ -40,7 +40,9 @@ def initialize(global_state, store, document, params)
 
       sig { override.returns(T.nilable(Interface::WorkspaceEdit)) }
       def perform
-        char_position = @document.create_scanner.find_char_position(@position)
+        char_position = @global_state.synchronize do
+          @document.create_scanner.find_char_position(@position)
+        end
 
         node_context = RubyDocument.locate(
           @document.parse_result.value,
diff --git a/lib/ruby_lsp/requests/show_syntax_tree.rb b/lib/ruby_lsp/requests/show_syntax_tree.rb
index 2852f90ea..62ad1f939 100644
--- a/lib/ruby_lsp/requests/show_syntax_tree.rb
+++ b/lib/ruby_lsp/requests/show_syntax_tree.rb
@@ -9,9 +9,12 @@ module Requests
     class ShowSyntaxTree < Request
       extend T::Sig
 
-      sig { params(document: RubyDocument, range: T.nilable(T::Hash[Symbol, T.untyped])).void }
-      def initialize(document, range)
+      sig do
+        params(global_state: GlobalState, document: RubyDocument, range: T.nilable(T::Hash[Symbol, T.untyped])).void
+      end
+      def initialize(global_state, document, range)
         super()
+        @global_state = global_state
         @document = document
         @range = range
         @tree = T.let(document.parse_result.value, Prism::ProgramNode)
@@ -32,9 +35,10 @@ def perform
       def ast_for_range
         range = T.must(@range)
 
-        scanner = @document.create_scanner
-        start_char = scanner.find_char_position(range[:start])
-        end_char = scanner.find_char_position(range[:end])
+        start_char, end_char = @global_state.synchronize do
+          scanner = @document.create_scanner
+          [scanner.find_char_position(range[:start]), scanner.find_char_position(range[:end])]
+        end
 
         queue = @tree.statements.body.dup
         found_nodes = []
diff --git a/lib/ruby_lsp/requests/signature_help.rb b/lib/ruby_lsp/requests/signature_help.rb
index ce8709d88..5ed6aad2c 100644
--- a/lib/ruby_lsp/requests/signature_help.rb
+++ b/lib/ruby_lsp/requests/signature_help.rb
@@ -36,7 +36,10 @@ def provider
       def initialize(document, global_state, position, context, dispatcher, sorbet_level) # rubocop:disable Metrics/ParameterLists
         super()
 
-        char_position = document.create_scanner.find_char_position(position)
+        char_position = global_state.synchronize do
+          document.create_scanner.find_char_position(position)
+        end
+
         delegate_request_if_needed!(global_state, document, char_position)
 
         node_context = RubyDocument.locate(
diff --git a/lib/ruby_lsp/server.rb b/lib/ruby_lsp/server.rb
index f975e065e..0e8e0e28e 100644
--- a/lib/ruby_lsp/server.rb
+++ b/lib/ruby_lsp/server.rb
@@ -107,7 +107,7 @@ def process_message(message)
           ),
         )
       when "$/cancelRequest"
-        @mutex.synchronize { @cancelled_requests << message[:params][:id] }
+        @global_state.synchronize { @cancelled_requests << message[:params][:id] }
       when nil
         process_response(message) if message[:result]
       end
@@ -377,7 +377,7 @@ def run_initialized
 
     sig { params(message: T::Hash[Symbol, T.untyped]).void }
     def text_document_did_open(message)
-      @mutex.synchronize do
+      @global_state.synchronize do
         text_document = message.dig(:params, :textDocument)
         language_id = case text_document[:languageId]
         when "erb", "eruby"
@@ -417,7 +417,7 @@ def text_document_did_open(message)
 
     sig { params(message: T::Hash[Symbol, T.untyped]).void }
     def text_document_did_close(message)
-      @mutex.synchronize do
+      @global_state.synchronize do
         uri = message.dig(:params, :textDocument, :uri)
         @store.delete(uri)
 
@@ -436,7 +436,7 @@ def text_document_did_change(message)
       params = message[:params]
       text_document = params[:textDocument]
 
-      @mutex.synchronize do
+      @global_state.synchronize do
         @store.push_edits(uri: text_document[:uri], edits: params[:contentChanges], version: text_document[:version])
       end
     end
@@ -755,7 +755,7 @@ def text_document_prepare_rename(message)
       send_message(
         Result.new(
           id: message[:id],
-          response: Requests::PrepareRename.new(document, params[:position]).perform,
+          response: Requests::PrepareRename.new(@global_state, document, params[:position]).perform,
         ),
       )
     end
@@ -1056,6 +1056,7 @@ def text_document_show_syntax_tree(message)
 
       response = {
         ast: Requests::ShowSyntaxTree.new(
+          @global_state,
           document,
           params[:range],
         ).perform,
diff --git a/test/requests/prepare_rename_expectations_test.rb b/test/requests/prepare_rename_expectations_test.rb
index 9922e1a07..896b98451 100644
--- a/test/requests/prepare_rename_expectations_test.rb
+++ b/test/requests/prepare_rename_expectations_test.rb
@@ -11,7 +11,8 @@ def run_expectations(source)
     position = @__params&.any? ? @__params[:position] : default_position
     uri = URI("file://#{@_path}")
     document = RubyLsp::RubyDocument.new(source: source, version: 1, uri: uri)
-    RubyLsp::Requests::PrepareRename.new(document, position).perform
+    global_state = RubyLsp::GlobalState.new
+    RubyLsp::Requests::PrepareRename.new(global_state, document, position).perform
   end
 
   private