Better cleanup...

Author: cmaglie

Diff for: handler/handler.go

@@ -51,6 +51,7 @@ type InoHandler struct {
 	clangdNotificationCount int64
 	progressHandler         *ProgressProxyHandler
 
+	closing                    chan bool
 	clangdStarted              *sync.Cond
 	dataMux                    sync.RWMutex
 	lspInitializeParams        *lsp.InitializeParams
@@ -119,6 +120,7 @@ func NewInoHandler(stdio io.ReadWriteCloser, board lsp.Board) *InoHandler {
 	handler := &InoHandler{
 		docs:                   map[string]*lsp.TextDocumentItem{},
 		inoDocsWithDiagnostics: map[lsp.DocumentURI]bool{},
+		closing:                make(chan bool),
 		config: lsp.BoardConfig{
 			SelectedBoard: board,
 		},
@@ -158,15 +160,27 @@ type FileData struct {
 	version    int
 }
 
-// StopClangd closes the connection to the clangd process.
-func (handler *InoHandler) StopClangd() {
-	handler.ClangdConn.Close()
-	handler.ClangdConn = nil
+// Close closes all the json-rpc connections.
+func (handler *InoHandler) Close() {
+	if handler.ClangdConn != nil {
+		handler.ClangdConn.Close()
+		handler.ClangdConn = nil
+	}
+	if handler.closing != nil {
+		close(handler.closing)
+		handler.closing = nil
+	}
+}
+
+// CloseNotify returns a channel that is closed when the InoHandler is closed
+func (handler *InoHandler) CloseNotify() <-chan bool {
+	return handler.closing
 }
 
 // CleanUp performs cleanup of the workspace and temp files create by the language server
 func (handler *InoHandler) CleanUp() {
 	if handler.buildPath != nil {
+		log.Printf("removing buildpath")
 		handler.buildPath.RemoveAll()
 		handler.buildPath = nil
 	}
@@ -489,12 +503,14 @@ func (handler *InoHandler) HandleMessageFromIDE(ctx context.Context, conn *jsonr
 		// Exit the process and trigger a restart by the client in case of a severe error
 		if err.Error() == "context deadline exceeded" {
 			log.Println(prefix + "Timeout exceeded while waiting for a reply from clangd.")
-			handler.exit()
+			log.Println(prefix + "Please restart the language server.")
+			handler.Close()
 		}
 		if strings.Contains(err.Error(), "non-added document") || strings.Contains(err.Error(), "non-added file") {
 			log.Printf(prefix + "The clangd process has lost track of the open document.")
 			log.Printf(prefix+"  %s", err)
-			handler.exit()
+			log.Println(prefix + "Please restart the language server.")
+			handler.Close()
 		}
 	}
 
@@ -505,12 +521,6 @@ func (handler *InoHandler) HandleMessageFromIDE(ctx context.Context, conn *jsonr
 	return result, err
 }
 
-func (handler *InoHandler) exit() {
-	log.Println("Please restart the language server.")
-	handler.StopClangd()
-	os.Exit(1)
-}
-
 func (handler *InoHandler) initializeWorkbench(ctx context.Context, params *lsp.InitializeParams) error {
 	currCppTextVersion := 0
 	if params != nil {
@@ -579,6 +589,11 @@ func (handler *InoHandler) initializeWorkbench(ctx context.Context, params *lsp.
 		handler.ClangdConn = jsonrpc2.NewConn(context.Background(), clangdStream, clangdHandler,
 			jsonrpc2.OnRecv(streams.JSONRPCConnLogOnRecv("IDE     LS <-- CL:")),
 			jsonrpc2.OnSend(streams.JSONRPCConnLogOnSend("IDE     LS --> CL:")))
+		go func() {
+			<-handler.ClangdConn.DisconnectNotify()
+			log.Printf("Lost connection with clangd!")
+			handler.Close()
+		}()
 
 		// Send initialization command to clangd
 		ctx, cancel := context.WithTimeout(context.Background(), time.Second)

Diff for: main.go

@@ -5,6 +5,7 @@ import (
 	"io"
 	"log"
 	"os"
+	"os/signal"
 
 	"github.com/arduino/go-paths-helper"
 	"github.com/bcmi-labs/arduino-language-server/handler"
@@ -53,7 +54,16 @@ func main() {
 	}
 
 	inoHandler := handler.NewInoHandler(stdio, initialBoard)
-	defer inoHandler.StopClangd()
-	defer inoHandler.CleanUp()
-	<-inoHandler.StdioConn.DisconnectNotify()
+
+	// Intercept kill signal
+	c := make(chan os.Signal, 2)
+	signal.Notify(c, os.Interrupt, os.Kill)
+
+	select {
+	case <-inoHandler.CloseNotify():
+	case <-c:
+		log.Println("INTERRUPTED")
+	}
+	inoHandler.CleanUp()
+	inoHandler.Close()
 }