Skip to content

Commit 2d33b3a

Browse files
committed
Fix Language Server not stopping when client disconnects
Previously when the client sent an `exit` message, the langauge server would attempt to close the underlying connection however it would hang, and at worst cause high CPU usage. This commit - Uses the callback architecture inside the TCP Server to close connection instead of trying to close the TCP Socker within the same execution context. This stops the possibility of a race condition which can cause unexpected blocks - Add a missing method `close_connection_after_writing` on the base connection class - Add more logging on TCP startup to describe whether the server will close after the client disconnects
1 parent 0df1fb9 commit 2d33b3a

File tree

3 files changed

+19
-6
lines changed

3 files changed

+19
-6
lines changed

server/lib/puppet-languageserver/json_rpc_handler.rb

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -55,7 +55,7 @@ def post_init
5555
end
5656

5757
def unbind
58-
PuppetLanguageServer::LogMessage('information','Client has disconnected to the language server')
58+
PuppetLanguageServer::LogMessage('information','Client has disconnected from the language server')
5959
end
6060

6161
def extract_headers(raw_header)

server/lib/puppet-languageserver/message_router.rb

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -126,7 +126,7 @@ def receive_notification(method, params)
126126
PuppetLanguageServer::LogMessage('information','Client has received initialization')
127127

128128
when 'exit'
129-
PuppetLanguageServer::LogMessage('information','Received exit notification. Shutting down.')
129+
PuppetLanguageServer::LogMessage('information','Received exit notification. Closing connection to client...')
130130
close_connection
131131

132132
when 'textDocument/didOpen'

server/lib/puppet-languageserver/simple_tcp_server.rb

Lines changed: 17 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,8 @@
66
module PuppetLanguageServer
77
class SimpleTCPServerConnection
88
attr_accessor :socket
9+
attr_accessor :simple_tcp_server
10+
911
# Methods to override
1012
def post_init
1113
# Override this to recieve events after a client is connected
@@ -29,14 +31,19 @@ def error?
2931

3032
# @api public
3133
def send_data(data)
32-
@socket.write(data)
33-
34+
socket.write(data)
3435
true
3536
end
3637

38+
# @api public
39+
def close_connection_after_writing
40+
socket.flush
41+
simple_tcp_server.remove_connection_async(socket)
42+
end
43+
3744
# @api public
3845
def close_connection
39-
socket.close
46+
simple_tcp_server.remove_connection_async(socket)
4047
end
4148
end
4249

@@ -100,6 +107,7 @@ def start(handler = PuppetLanguageServer::SimpleTCPServerConnection, connection_
100107
kill_timer = connection_options[:connection_timeout]
101108
kill_timer = -1 if kill_timer.nil? || kill_timer < 1
102109
log("Will stop the server in #{connection_options[:connection_timeout]} seconds if no connection is made.") if kill_timer > 0
110+
log('Will stop the server when client disconnects') if !@server_options[:stop_on_client_exit].nil? && @server_options[:stop_on_client_exit]
103111

104112
(begin
105113
sleep(1)
@@ -241,6 +249,11 @@ def stop_services
241249
S_LOCKER.synchronize {SERVICES.each {|s, p| (s.close rescue true); log("Stopped listening on #{p[:hostname]}:#{p[:port]}") }; SERVICES.clear }
242250
end
243251

252+
# @api public
253+
def remove_connection_async(io)
254+
callback(self, :remove_connection, io)
255+
end
256+
244257
#####################
245258
# IO - Active connections handling
246259

@@ -252,7 +265,7 @@ def stop_connections
252265
def add_connection(io, service_object)
253266
handler = @handler_klass.new(@handler_start_options)
254267
handler.socket = io
255-
268+
handler.simple_tcp_server = self
256269
C_LOCKER.synchronize {IO_CONNECTION_DIC[io] = { :handler => handler, :service => service_object} } if io
257270
callback(handler, :post_init)
258271
end

0 commit comments

Comments
 (0)