internal/lsp: add shutdown handling

We correcly cancel all background tasks and drop all active views when
the server is asked to shut down now.
This was mostly to support the command line being able to exit cleanly

Change-Id: Iff9f5ab51572aad5e3245dc01aa87b00dcd47963
Reviewed-on: https://go-review.googlesource.com/c/tools/+/174940
Run-TryBot: Ian Cottrell <iancottrell@google.com>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Rebecca Stambler <rstambler@golang.org>
diff --git a/internal/lsp/cache/view.go b/internal/lsp/cache/view.go
index 4329821..8117065 100644
--- a/internal/lsp/cache/view.go
+++ b/internal/lsp/cache/view.go
@@ -138,6 +138,15 @@
 	v.config.Env = env
 }
 
+func (v *view) Shutdown(context.Context) {
+	v.mu.Lock()
+	defer v.mu.Unlock()
+	if v.cancel != nil {
+		v.cancel()
+		v.cancel = nil
+	}
+}
+
 func (v *view) BackgroundContext() context.Context {
 	v.mu.Lock()
 	defer v.mu.Unlock()
diff --git a/internal/lsp/cmd/check.go b/internal/lsp/cmd/check.go
index 3231828..59f0599 100644
--- a/internal/lsp/cmd/check.go
+++ b/internal/lsp/cmd/check.go
@@ -45,6 +45,7 @@
 	if err != nil {
 		return err
 	}
+	defer conn.terminate(ctx)
 	for _, arg := range args {
 		uri := span.FileURI(arg)
 		file := conn.AddFile(ctx, uri)
diff --git a/internal/lsp/cmd/cmd.go b/internal/lsp/cmd/cmd.go
index 1059f04..45ce925 100644
--- a/internal/lsp/cmd/cmd.go
+++ b/internal/lsp/cmd/cmd.go
@@ -340,3 +340,14 @@
 	}
 	return file
 }
+
+func (c *connection) terminate(ctx context.Context) {
+	if c.Client.app.Remote == "internal" {
+		// internal connections need to be left alive for the next test
+		return
+	}
+	//TODO: do we need to handle errors on these calls?
+	c.Shutdown(ctx)
+	//TODO: right now calling exit terminates the process, we should rethink that
+	//server.Exit(ctx)
+}
diff --git a/internal/lsp/cmd/definition.go b/internal/lsp/cmd/definition.go
index d093ab7..4c78327 100644
--- a/internal/lsp/cmd/definition.go
+++ b/internal/lsp/cmd/definition.go
@@ -63,6 +63,7 @@
 	if err != nil {
 		return err
 	}
+	defer conn.terminate(ctx)
 	from := span.Parse(args[0])
 	file := conn.AddFile(ctx, from.URI())
 	if file.err != nil {
diff --git a/internal/lsp/cmd/format.go b/internal/lsp/cmd/format.go
index 0f6542e..4075ec4 100644
--- a/internal/lsp/cmd/format.go
+++ b/internal/lsp/cmd/format.go
@@ -55,6 +55,7 @@
 	if err != nil {
 		return err
 	}
+	defer conn.terminate(ctx)
 	for _, arg := range args {
 		spn := span.Parse(arg)
 		file := conn.AddFile(ctx, spn.URI())
diff --git a/internal/lsp/general.go b/internal/lsp/general.go
index dbc3703..38f9b9c 100644
--- a/internal/lsp/general.go
+++ b/internal/lsp/general.go
@@ -189,12 +189,19 @@
 }
 
 func (s *Server) shutdown(ctx context.Context) error {
-	// TODO(rstambler): Cancel contexts here?
 	s.initializedMu.Lock()
 	defer s.initializedMu.Unlock()
 	if !s.isInitialized {
 		return jsonrpc2.NewErrorf(jsonrpc2.CodeInvalidRequest, "server not initialized")
 	}
+	// drop all the active views
+	s.viewMu.Lock()
+	defer s.viewMu.Unlock()
+	for _, v := range s.views {
+		v.Shutdown(ctx)
+	}
+	s.views = nil
+	s.viewMap = nil
 	s.isInitialized = false
 	return nil
 }
diff --git a/internal/lsp/source/view.go b/internal/lsp/source/view.go
index 44796a7..91c82c0 100644
--- a/internal/lsp/source/view.go
+++ b/internal/lsp/source/view.go
@@ -32,6 +32,7 @@
 	BackgroundContext() context.Context
 	Config() packages.Config
 	SetEnv([]string)
+	Shutdown(ctx context.Context)
 }
 
 // File represents a Go source file that has been type-checked. It is the input
diff --git a/internal/lsp/workspace.go b/internal/lsp/workspace.go
index 44674b3..938450a 100644
--- a/internal/lsp/workspace.go
+++ b/internal/lsp/workspace.go
@@ -77,7 +77,7 @@
 			s.views[i] = s.views[len(s.views)-1]
 			s.views[len(s.views)-1] = nil
 			s.views = s.views[:len(s.views)-1]
-			//TODO: shutdown the view in here
+			view.Shutdown(ctx)
 			return nil
 		}
 	}