Merge pull request #33 from glennsarti/fix-client-disconnection

Fix Language Server not stopping when client disconnects

Author: jpogran

server/lib/puppet-languageserver/json_rpc_handler.rb

@@ -55,7 +55,7 @@ def post_init
     end
 
     def unbind
-      PuppetLanguageServer::LogMessage('information','Client has disconnected to the language server')
+      PuppetLanguageServer::LogMessage('information','Client has disconnected from the language server')
     end
 
     def extract_headers(raw_header)

server/lib/puppet-languageserver/message_router.rb

@@ -126,7 +126,7 @@ def receive_notification(method, params)
           PuppetLanguageServer::LogMessage('information','Client has received initialization')
 
         when 'exit'
-          PuppetLanguageServer::LogMessage('information','Received exit notification.  Shutting down.')
+          PuppetLanguageServer::LogMessage('information','Received exit notification.  Closing connection to client...')
           close_connection
 
         when 'textDocument/didOpen'

server/lib/puppet-languageserver/simple_tcp_server.rb

@@ -6,6 +6,8 @@
 module PuppetLanguageServer
   class SimpleTCPServerConnection
     attr_accessor :socket
+    attr_accessor :simple_tcp_server
+
     # Methods to override
     def post_init
       # Override this to recieve events after a client is connected
@@ -29,14 +31,19 @@ def error?
 
     # @api public
     def send_data(data)
-      @socket.write(data)
-
+      socket.write(data)
       true
     end
 
+    # @api public
+    def close_connection_after_writing
+      socket.flush
+      simple_tcp_server.remove_connection_async(socket)
+    end
+
     # @api public
     def close_connection
-      socket.close
+      simple_tcp_server.remove_connection_async(socket)
     end
   end
 
@@ -100,6 +107,7 @@ def start(handler = PuppetLanguageServer::SimpleTCPServerConnection, connection_
       kill_timer = connection_options[:connection_timeout]
       kill_timer = -1 if kill_timer.nil? || kill_timer < 1
       log("Will stop the server in #{connection_options[:connection_timeout]} seconds if no connection is made.") if kill_timer > 0
+      log('Will stop the server when client disconnects') if !@server_options[:stop_on_client_exit].nil? && @server_options[:stop_on_client_exit]
 
       (begin
         sleep(1)
@@ -241,6 +249,11 @@ def stop_services
       S_LOCKER.synchronize {SERVICES.each {|s, p| (s.close rescue true); log("Stopped listening on #{p[:hostname]}:#{p[:port]}") }; SERVICES.clear }
     end
 
+    # @api public
+    def remove_connection_async(io)
+      callback(self, :remove_connection, io)
+    end
+
     #####################
     # IO - Active connections handling
 
@@ -252,7 +265,7 @@ def stop_connections
     def add_connection(io, service_object)
       handler = @handler_klass.new(@handler_start_options)
       handler.socket = io
-
+      handler.simple_tcp_server = self
       C_LOCKER.synchronize {IO_CONNECTION_DIC[io] = { :handler => handler, :service => service_object} } if io
       callback(handler, :post_init)
     end