From 7147197327d0bbfeed0b089b7c4df2552466dae8 Mon Sep 17 00:00:00 2001
From: Ian Cottrell <iancottrell@google.com>
Date: Fri, 5 Jun 2020 09:17:05 -0400
Subject: [PATCH] internal/lsp: now connection shutdown works, use it

This removes a TODO in the tests and has them check for shutdown
correctly.
This also required adding a Close to editor that does the Shutdown/Exit
sequence rather than just Shutdown.

Change-Id: I6433b76eb2dd1fe40b2332685fdf25ebc4fd4577
Reviewed-on: https://go-review.googlesource.com/c/tools/+/236717
Run-TryBot: Ian Cottrell <iancottrell@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Robert Findley <rfindley@google.com>
---
 internal/lsp/fake/editor.go        | 26 +++++++++++++++++++++++---
 internal/lsp/lsprpc/lsprpc_test.go | 27 ++++++++++++++++++++++-----
 internal/lsp/regtest/wrappers.go   |  3 +--
 3 files changed, 46 insertions(+), 10 deletions(-)

diff --git a/internal/lsp/fake/editor.go b/internal/lsp/fake/editor.go
index f1271c5298..50601313d3 100644
--- a/internal/lsp/fake/editor.go
+++ b/internal/lsp/fake/editor.go
@@ -24,9 +24,10 @@ type Editor struct {
 
 	// Server, client, and sandbox are concurrency safe and written only
 	// at construction time, so do not require synchronization.
-	Server  protocol.Server
-	client  *Client
-	sandbox *Sandbox
+	Server     protocol.Server
+	serverConn jsonrpc2.Conn
+	client     *Client
+	sandbox    *Sandbox
 
 	// Since this editor is intended just for testing, we use very coarse
 	// locking.
@@ -81,6 +82,7 @@ func NewEditor(ws *Sandbox, config EditorConfig) *Editor {
 // It returns the editor, so that it may be called as follows:
 //   editor, err := NewEditor(s).Connect(ctx, conn)
 func (e *Editor) Connect(ctx context.Context, conn jsonrpc2.Conn, hooks ClientHooks) (*Editor, error) {
+	e.serverConn = conn
 	e.Server = protocol.ServerDispatcher(conn)
 	e.client = &Client{editor: e, hooks: hooks}
 	conn.Go(ctx,
@@ -116,6 +118,24 @@ func (e *Editor) Exit(ctx context.Context) error {
 	return nil
 }
 
+// Close issues the shutdown and exit sequence an editor should.
+func (e *Editor) Close(ctx context.Context) error {
+	if err := e.Shutdown(ctx); err != nil {
+		return err
+	}
+	if err := e.Exit(ctx); err != nil {
+		return err
+	}
+	// called close on the editor should result in the connection closing
+	select {
+	case <-e.serverConn.Done():
+		// connection closed itself
+		return nil
+	case <-ctx.Done():
+		return fmt.Errorf("connection not closed: %w", ctx.Err())
+	}
+}
+
 // Client returns the LSP client for this editor.
 func (e *Editor) Client() *Client {
 	return e.client
diff --git a/internal/lsp/lsprpc/lsprpc_test.go b/internal/lsp/lsprpc/lsprpc_test.go
index 930fb6a37c..01e57cea88 100644
--- a/internal/lsp/lsprpc/lsprpc_test.go
+++ b/internal/lsp/lsprpc/lsprpc_test.go
@@ -219,13 +219,13 @@ func TestDebugInfoLifecycle(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer ed1.Shutdown(clientCtx)
+	defer ed1.Close(clientCtx)
 	conn2 := tsBackend.Connect(baseCtx)
 	ed2, err := fake.NewEditor(sb, fake.EditorConfig{}).Connect(baseCtx, conn2, fake.ClientHooks{})
 	if err != nil {
 		t.Fatal(err)
 	}
-	defer ed2.Shutdown(baseCtx)
+	defer ed2.Close(baseCtx)
 
 	serverDebug := debug.GetInstance(serverCtx)
 	if got, want := len(serverDebug.State.Clients()), 2; got != want {
@@ -240,12 +240,29 @@ func TestDebugInfoLifecycle(t *testing.T) {
 	}
 	// Close one of the connections to verify that the client and session were
 	// dropped.
-	if err := ed1.Shutdown(clientCtx); err != nil {
+	if err := ed1.Close(clientCtx); err != nil {
 		t.Fatal(err)
 	}
+	/*TODO: at this point we have verified the editor is closed
+	However there is no way currently to wait for all associated go routines to
+	go away, and we need to wait for those to trigger the client drop
+	for now we just give it a little bit of time, but we need to fix this
+	in a principled way
+	*/
+	start := time.Now()
+	delay := time.Millisecond
+	const maxWait = time.Second
+	for len(serverDebug.State.Clients()) > 1 {
+		if time.Since(start) > maxWait {
+			break
+		}
+		time.Sleep(delay)
+		delay *= 2
+	}
+	if got, want := len(serverDebug.State.Clients()), 1; got != want {
+		t.Errorf("len(server:Clients) = %d, want %d", got, want)
+	}
 	if got, want := len(serverDebug.State.Sessions()), 1; got != want {
 		t.Errorf("len(server:Sessions()) = %d, want %d", got, want)
 	}
-	// TODO(rfindley): once disconnection works, assert that len(Clients) == 1
-	// (as of writing, it is still 2)
 }
diff --git a/internal/lsp/regtest/wrappers.go b/internal/lsp/regtest/wrappers.go
index aaf9199089..276f5b19cc 100644
--- a/internal/lsp/regtest/wrappers.go
+++ b/internal/lsp/regtest/wrappers.go
@@ -165,8 +165,7 @@ func checkIsFatal(t *testing.T, err error) {
 // CloseEditor shuts down the editor, calling t.Fatal on any error.
 func (e *Env) CloseEditor() {
 	e.T.Helper()
-	checkIsFatal(e.T, e.Editor.Shutdown(e.Ctx))
-	checkIsFatal(e.T, e.Editor.Exit(e.Ctx))
+	checkIsFatal(e.T, e.Editor.Close(e.Ctx))
 }
 
 // RunGenerate runs go:generate on the given dir, calling t.Fatal on any error.