From b270166f7ab5fcd4aca77fda83ae99a1c678be42 Mon Sep 17 00:00:00 2001
From: Kevin Gillette <kgillette@twilio.com>
Date: Sun, 16 Feb 2025 19:11:06 -0700
Subject: [PATCH] handle 'exit' message

At least with the helix editor, prior to this change,
golangci-lint-langserver will handle the 'shutdown' message,
but will then stall, waiting for the invoking client (helix)
to terminate the LSP process, while helix is waiting for the LSP process
to self-terminate upon receiving the 'exit' message.
Eventually, helix will terminate the LSP process
after a ~10 second grace period has passed,
though this is rather annoying to wait for.

With the change, the 'exit' command will cause
the jsonrpc2 connection to close,
which in turn will lead to the LSP process exiting quickly.
---
 handler.go | 47 ++++++++++++++++++++++++++++++-----------------
 main.go    | 31 ++++++++++++++++++++++++-------
 2 files changed, 54 insertions(+), 24 deletions(-)

diff --git a/handler.go b/handler.go
index 99213e9..f2e5ee2 100644
--- a/handler.go
+++ b/handler.go
@@ -11,12 +11,20 @@ import (
 	"github.com/sourcegraph/jsonrpc2"
 )
 
-func NewHandler(logger logger, noLinterName bool) jsonrpc2.Handler {
+type HandlerConfig struct {
+	Logger       logger
+	NoLinterName bool
+	Close        func()
+}
+
+func NewHandler(cfg *HandlerConfig) jsonrpc2.Handler {
 	handler := &langHandler{
-		logger:       logger,
+		logger:       cfg.Logger,
 		request:      make(chan DocumentURI),
-		noLinterName: noLinterName,
+		noLinterName: cfg.NoLinterName,
+		close:        cfg.Close,
 	}
+
 	go handler.linter()
 
 	return jsonrpc2.HandlerWithError(handler.handle)
@@ -28,6 +36,7 @@ type langHandler struct {
 	request      chan DocumentURI
 	command      []string
 	noLinterName bool
+	close        func()
 
 	rootURI string
 	rootDir string
@@ -110,13 +119,6 @@ func (h *langHandler) lint(uri DocumentURI) ([]Diagnostic, error) {
 	return diagnostics, nil
 }
 
-func max(a, b int) int {
-	if a > b {
-		return a
-	}
-	return b
-}
-
 func (h *langHandler) diagnosticMessage(issue *Issue) string {
 	if h.noLinterName {
 		return issue.Text
@@ -126,12 +128,7 @@ func (h *langHandler) diagnosticMessage(issue *Issue) string {
 }
 
 func (h *langHandler) linter() {
-	for {
-		uri, ok := <-h.request
-		if !ok {
-			break
-		}
-
+	for uri := range h.request {
 		diagnostics, err := h.lint(uri)
 		if err != nil {
 			h.logger.Printf("%s", err)
@@ -161,6 +158,8 @@ func (h *langHandler) handle(ctx context.Context, conn *jsonrpc2.Conn, req *json
 		return
 	case "shutdown":
 		return h.handleShutdown(ctx, conn, req)
+	case "exit":
+		return h.handleExit(ctx, conn, req)
 	case "textDocument/didOpen":
 		return h.handleTextDocumentDidOpen(ctx, conn, req)
 	case "textDocument/didClose":
@@ -198,8 +197,22 @@ func (h *langHandler) handleInitialize(_ context.Context, conn *jsonrpc2.Conn, r
 	}, nil
 }
 
+func (h *langHandler) shutdown() {
+	if h.request != nil {
+		close(h.request)
+		h.request = nil
+	}
+}
+
 func (h *langHandler) handleShutdown(_ context.Context, _ *jsonrpc2.Conn, _ *jsonrpc2.Request) (result interface{}, err error) {
-	close(h.request)
+	h.shutdown()
+
+	return nil, nil
+}
+
+func (h *langHandler) handleExit(_ context.Context, _ *jsonrpc2.Conn, _ *jsonrpc2.Request) (result interface{}, err error) {
+	h.shutdown()
+	h.close()
 
 	return nil, nil
 }
diff --git a/main.go b/main.go
index 7848131..cb57d4c 100644
--- a/main.go
+++ b/main.go
@@ -19,20 +19,37 @@ func main() {
 
 	logger := newStdLogger(*debug)
 
-	handler := NewHandler(logger, *noLinterName)
+	ctx := context.Background()
+	ctx, cancel := context.WithCancel(ctx)
+	defer cancel()
 
-	var connOpt []jsonrpc2.ConnOpt
+	handler := NewHandler(&HandlerConfig{
+		Logger:       logger,
+		NoLinterName: *noLinterName,
+		Close:        cancel,
+	})
 
-	logger.Printf("golangci-lint-langserver: connections opened")
+	var connOpt []jsonrpc2.ConnOpt
+	if *debug {
+		connOpt = append(connOpt, jsonrpc2.LogMessages(logger))
+	}
 
-	<-jsonrpc2.NewConn(
-		context.Background(),
+	conn := jsonrpc2.NewConn(
+		ctx,
 		jsonrpc2.NewBufferedStream(stdrwc{}, jsonrpc2.VSCodeObjectCodec{}),
 		handler,
 		connOpt...,
-	).DisconnectNotify()
+	)
 
-	logger.Printf("golangci-lint-langserver: connections closed")
+	defer conn.Close()
+
+	logger.Printf("golangci-lint-langserver: connection opened")
+	defer logger.Printf("golangci-lint-langserver: connection closed")
+
+	select {
+	case <-ctx.Done():
+	case <-conn.DisconnectNotify():
+	}
 }
 
 type stdrwc struct{}