Increase request concurrency & make it configurable

internal/cmd/serve_command.go

@@ -26,13 +26,14 @@ type ServeCommand struct {
 	Version string
 
 	// flags
-	port          int
-	logFilePath   string
-	tfExecPath    string
-	tfExecLogPath string
-	tfExecTimeout string
-	cpuProfile    string
-	memProfile    string
+	port           int
+	logFilePath    string
+	tfExecPath     string
+	tfExecLogPath  string
+	tfExecTimeout  string
+	cpuProfile     string
+	memProfile     string
+	reqConcurrency int
 }
 
 func (c *ServeCommand) flags() *flag.FlagSet {
@@ -52,6 +53,8 @@ func (c *ServeCommand) flags() *flag.FlagSet {
 	fs.StringVar(&c.memProfile, "memprofile", "", "file into which to write memory profile (if not empty)"+
 		" with support for variables (e.g. Timestamp, Pid, Ppid) via Go template"+
 		" syntax {{.VarName}}")
+	fs.IntVar(&c.reqConcurrency, "req-concurrency", 0, fmt.Sprintf("number of RPC requests to process concurrently,"+
+		" defaults to %d, concurrency lower than 2 is not recommended", langserver.DefaultConcurrency()))
 
 	fs.Usage = func() { c.Ui.Error(c.Help()) }
 
@@ -139,6 +142,11 @@ func (c *ServeCommand) Run(args []string) int {
 		logger.Printf("Terraform exec path set to %q", path)
 	}
 
+	if c.reqConcurrency != 0 {
+		ctx = langserver.WithRequestConcurrency(ctx, c.reqConcurrency)
+		logger.Printf("Custom request concurrency set to %d", c.reqConcurrency)
+	}
+
 	logger.Printf("Starting terraform-ls %s", c.Version)
 
 	ctx = lsctx.WithLanguageServerVersion(ctx, c.Version)

internal/langserver/langserver.go

@@ -8,6 +8,7 @@ import (
 	"log"
 	"net"
 	"os"
+	"runtime"
 
 	"github.com/creachadair/jrpc2"
 	"github.com/creachadair/jrpc2/channel"
@@ -22,13 +23,17 @@ type langServer struct {
 	newSession session.SessionFactory
 }
 
+type ctxReqConcurrency struct{}
+
 func NewLangServer(srvCtx context.Context, sf session.SessionFactory) *langServer {
-	opts := &jrpc2.ServerOptions{
-		AllowPush: true,
+	concurrency, ok := requestConcurrencyFromCtx(srvCtx)
+	if !ok {
+		concurrency = DefaultConcurrency()
+	}
 
-		// Disable concurrency to avoid race conditions
-		// between requests concerning the same document
-		Concurrency: 1,
+	opts := &jrpc2.ServerOptions{
+		AllowPush:   true,
+		Concurrency: concurrency,
 	}
 
 	return &langServer{
@@ -39,6 +44,26 @@ func NewLangServer(srvCtx context.Context, sf session.SessionFactory) *langServe
 	}
 }
 
+func WithRequestConcurrency(parent context.Context, concurrency int) context.Context {
+	return context.WithValue(parent, ctxReqConcurrency{}, concurrency)
+}
+
+func requestConcurrencyFromCtx(ctx context.Context) (int, bool) {
+	c, ok := ctx.Value(ctxReqConcurrency{}).(int)
+	return c, ok
+}
+
+func DefaultConcurrency() int {
+	cpu := runtime.NumCPU()
+	// Cap concurrency on powerful machines
+	// to leave some capacity for module ops
+	// and other application
+	if cpu >= 4 {
+		return cpu / 2
+	}
+	return cpu
+}
+
 func (ls *langServer) SetLogger(logger *log.Logger) {
 	ls.srvOptions.Logger = logger
 	ls.srvOptions.RPCLog = &rpcLogger{logger}
@@ -66,7 +91,8 @@ func (ls *langServer) StartAndWait(reader io.Reader, writer io.WriteCloser) erro
 	if err != nil {
 		return err
 	}
-	ls.logger.Printf("Starting server (pid %d) ...", os.Getpid())
+	ls.logger.Printf("Starting server (pid %d; concurrency: %d) ...",
+		os.Getpid(), ls.srvOptions.Concurrency)
 
 	// Wrap waiter with a context so that we can cancel it here
 	// after the service is cancelled (and srv.Wait returns)
@@ -87,8 +113,8 @@ func (ls *langServer) StartAndWait(reader io.Reader, writer io.WriteCloser) erro
 }
 
 func (ls *langServer) StartTCP(address string) error {
-	ls.logger.Printf("Starting TCP server (pid %d) at %q ...",
-		os.Getpid(), address)
+	ls.logger.Printf("Starting TCP server (pid %d; concurrency: %d) at %q ...",
+		os.Getpid(), ls.srvOptions.Concurrency, address)
 	lst, err := net.Listen("tcp", address)
 	if err != nil {
 		return fmt.Errorf("TCP Server failed to start: %s", err)