Merge pull request sourcegraph#101 from sourcegraph/sg/context-timeout

slimsag · web-flow · commit eeb1a981b2bf · 2017-01-21T22:29:19.000-07:00
langserver: add 15s timeout to workspace/xreferences
diff --git a/langserver/workspace_refs.go b/langserver/workspace_refs.go
@@ -13,6 +13,7 @@ import (
 	"sort"
 	"strings"
 	"sync"
+	"time"
 
 	"golang.org/x/tools/go/loader"
 
@@ -23,7 +24,17 @@ import (
 	"github.com/sourcegraph/jsonrpc2"
 )
 
+// workspaceReferencesTimeout is the timeout used for workspace/xreferences
+// calls.
+const workspaceReferencesTimeout = 15 * time.Second
+
 func (h *LangHandler) handleWorkspaceReferences(ctx context.Context, conn JSONRPC2Conn, req *jsonrpc2.Request, params lspext.WorkspaceReferencesParams) ([]referenceInformation, error) {
+	// TODO: Add support for the cancelRequest LSP method instead of using
+	// hard-coded timeouts like this here.
+	//
+	// See: https://github.com/Microsoft/language-server-protocol/blob/master/protocol.md#cancelRequest
+	ctx, cancel := context.WithTimeout(ctx, workspaceReferencesTimeout)
+	defer cancel()
 	rootPath := h.FilePath(h.init.RootPath)
 	bctx := h.BuildContext(ctx)
 
@@ -101,13 +112,37 @@ func (h *LangHandler) handleWorkspaceReferences(ctx context.Context, conn JSONRP
 		}()
 	}
 
-	_, err := h.workspaceRefsTypecheck(ctx, bctx, conn, fset, pkgs, afterTypeCheck)
-	if err != nil {
-		return nil, err
-	}
+	// workspaceRefsTypecheck is ran inside it's own goroutine because it can
+	// block for longer than our context deadline.
+	var err error
+	done := make(chan struct{})
+	go func() {
+		// Prevent any uncaught panics from taking the entire server down.
+		defer func() {
+			if r := recover(); r != nil {
+				// Same as net/http
+				const size = 64 << 10
+				buf := make([]byte, size)
+				buf = buf[:runtime.Stack(buf, false)]
+				log.Printf("ignoring panic serving %v for pkgs %v: %v\n%s", req.Method, pkgs, r, buf)
+				return
+			}
+		}()
 
-	// Wait for all worker goroutines to complete.
-	wg.Wait()
+		_, err = h.workspaceRefsTypecheck(ctx, bctx, conn, fset, pkgs, afterTypeCheck)
+
+		// Wait for all worker goroutines to complete.
+		wg.Wait()
+		close(done)
+	}()
+	select {
+	case <-done:
+		if err != nil {
+			return nil, err
+		}
+	case <-ctx.Done():
+		return nil, ctx.Err()
+	}
 
 	sort.Sort(&results) // sort to provide consistent results
 	return results.results, nil
@@ -149,7 +184,12 @@ func (h *LangHandler) workspaceRefsTypecheck(ctx context.Context, bctx *build.Co
 			}
 			return bpkg, nil
 		},
-		AfterTypeCheck: afterTypeCheck,
+		AfterTypeCheck: func(pkg *loader.PackageInfo, files []*ast.File) {
+			if err := ctx.Err(); err != nil {
+				return
+			}
+			afterTypeCheck(pkg, files)
+		},
 	}
 	for _, path := range pkgs {
 		conf.Import(path)
@@ -161,6 +201,10 @@ func (h *LangHandler) workspaceRefsTypecheck(ctx context.Context, bctx *build.Co
 		return nil, err
 	}
 
+	if err := ctx.Err(); err != nil {
+		return nil, err
+	}
+
 	// Publish typechecking error diagnostics.
 	diags, err := errsToDiagnostics(typeErrs, prog)
 	if err != nil {
@@ -179,6 +223,9 @@ func (h *LangHandler) workspaceRefsTypecheck(ctx context.Context, bctx *build.Co
 // workspaceRefsFromPkg collects all the references made to dependencies from
 // the specified package and returns the results.
 func (h *LangHandler) workspaceRefsFromPkg(ctx context.Context, bctx *build.Context, conn JSONRPC2Conn, params lspext.WorkspaceReferencesParams, fs *token.FileSet, pkg *loader.PackageInfo, rootPath string, results *refResultSorter) (err error) {
+	if err := ctx.Err(); err != nil {
+		return err
+	}
 	span, ctx := opentracing.StartSpanFromContext(ctx, "workspaceRefsFromPkg")
 	defer func() {
 		if err != nil {
