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>

internal/lsp/cache/view.go

@@ -138,6 +138,15 @@ func (v *view) SetEnv(env []string) {
 	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()

internal/lsp/cmd/check.go

@@ -45,6 +45,7 @@ func (c *check) Run(ctx context.Context, args ...string) error {
 	if err != nil {
 		return err
 	}
+	defer conn.terminate(ctx)
 	for _, arg := range args {
 		uri := span.FileURI(arg)
 		file := conn.AddFile(ctx, uri)

internal/lsp/cmd/cmd.go

@@ -340,3 +340,14 @@ func (c *connection) AddFile(ctx context.Context, uri span.URI) *cmdFile {
 	}
 	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)
+}

internal/lsp/cmd/definition.go

@@ -63,6 +63,7 @@ func (d *definition) Run(ctx context.Context, args ...string) error {
 	if err != nil {
 		return err
 	}
+	defer conn.terminate(ctx)
 	from := span.Parse(args[0])
 	file := conn.AddFile(ctx, from.URI())
 	if file.err != nil {

internal/lsp/cmd/format.go

@@ -55,6 +55,7 @@ func (f *format) Run(ctx context.Context, args ...string) error {
 	if err != nil {
 		return err
 	}
+	defer conn.terminate(ctx)
 	for _, arg := range args {
 		spn := span.Parse(arg)
 		file := conn.AddFile(ctx, spn.URI())

internal/lsp/general.go

@@ -189,12 +189,19 @@ func applyEnv(env []string, k string, v interface{}) []string {
 }
 
 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
 }

internal/lsp/source/view.go

@@ -32,6 +32,7 @@ type View interface {
 	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

internal/lsp/workspace.go

@@ -77,7 +77,7 @@ func (s *Server) removeView(ctx context.Context, name string, uri span.URI) erro
 			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
 		}
 	}