internal/lsp: fix context cancellation

Change-Id: Ib932722b0e66910b8cd23031e756f16099ee82dc
Reviewed-on: https://go-review.googlesource.com/c/tools/+/214581
Run-TryBot: Rebecca Stambler <rstambler@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Heschi Kreinick <heschi@google.com>
diff --git a/internal/lsp/cache/load.go b/internal/lsp/cache/load.go
index c3357c5..e824b06 100644
--- a/internal/lsp/cache/load.go
+++ b/internal/lsp/cache/load.go
@@ -66,10 +66,11 @@
 	cfg := s.view.Config(ctx)
 	pkgs, err := packages.Load(cfg, query)
 
-	// If the context was canceled, return early.
-	// Otherwise, we might be type-checking an incomplete result.
-	if err == context.Canceled {
-		return nil, err
+	// If the context was canceled, return early. Otherwise, we might be
+	// type-checking an incomplete result. Check the context directly,
+	// because go/packages adds extra information to the error.
+	if ctx.Err() != nil {
+		return nil, ctx.Err()
 	}
 
 	log.Print(ctx, "go/packages.Load", tag.Of("query", query), tag.Of("packages", len(pkgs)))
diff --git a/internal/lsp/source/diagnostics.go b/internal/lsp/source/diagnostics.go
index 9dec184..6929f9b 100644
--- a/internal/lsp/source/diagnostics.go
+++ b/internal/lsp/source/diagnostics.go
@@ -117,6 +117,9 @@
 	// Updates to the diagnostics for this package may need to be propagated.
 	reverseDeps, err := snapshot.GetReverseDependencies(ctx, pkg.ID())
 	if err != nil {
+		if err == context.Canceled {
+			return nil, warningMsg, err
+		}
 		log.Error(ctx, "no reverse dependencies", err)
 		return reports, warningMsg, nil
 	}