Evict cache entries when the connection closes (#359)

Author: chrismwendt

Gopkg.lock

@@ -112,7 +112,7 @@ func BenchmarkIntegration(b *testing.B) {
 		label := strings.Replace(root.Host+root.Path, "/", "-", -1)
 
 		b.Run(label, func(b *testing.B) {
-			fs, err := gobuildserver.RemoteFS(context.Background(), lspext.InitializeParams{OriginalRootURI: rootURI})
+			fs, _, err := gobuildserver.RemoteFS(context.Background(), lspext.InitializeParams{OriginalRootURI: rootURI})
 			if err != nil {
 				b.Fatal(err)
 			}
@@ -232,11 +232,11 @@ func BenchmarkIntegrationShared(b *testing.B) {
 	ctx := context.Background()
 	for label, test := range tests {
 		b.Run(label, func(b *testing.B) {
-			oldfs, err := gobuildserver.RemoteFS(context.Background(), lspext.InitializeParams{OriginalRootURI: lsp.DocumentURI(test.oldRootURI)})
+			oldfs, _, err := gobuildserver.RemoteFS(context.Background(), lspext.InitializeParams{OriginalRootURI: lsp.DocumentURI(test.oldRootURI)})
 			if err != nil {
 				b.Fatal(err)
 			}
-			fs, err := gobuildserver.RemoteFS(context.Background(), lspext.InitializeParams{OriginalRootURI: lsp.DocumentURI(test.rootURI)})
+			fs, _, err := gobuildserver.RemoteFS(context.Background(), lspext.InitializeParams{OriginalRootURI: lsp.DocumentURI(test.rootURI)})
 			if err != nil {
 				b.Fatal(err)
 			}

buildserver/bench_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"log"
 	"net/http"
 	"net/url"
@@ -17,6 +18,8 @@ import (
 	"sync"
 	"time"
 
+	multierror "github.com/hashicorp/go-multierror"
+
 	"github.com/die-net/lrucache"
 	"github.com/gregjones/httpcache"
 	opentracing "github.com/opentracing/opentracing-go"
@@ -44,7 +47,7 @@ var Debug = true
 // and mapping local file system paths to logical URIs (e.g.,
 // /goroot/src/fmt/print.go ->
 // git://github.com/golang/go?go1.7.1#src/fmt/print.go).
-func NewHandler(defaultCfg langserver.Config) jsonrpc2.Handler {
+func NewHandler(defaultCfg langserver.Config) *BuildHandler {
 	if defaultCfg.MaxParallelism <= 0 {
 		panic(fmt.Sprintf("langserver.Config.MaxParallelism must be at least 1 (got %d)", defaultCfg.MaxParallelism))
 	}
@@ -57,7 +60,7 @@ func NewHandler(defaultCfg langserver.Config) jsonrpc2.Handler {
 		},
 	}
 	shared.FindPackage = h.findPackageCached
-	return jsonrpc2.HandlerWithError(h.handle)
+	return h
 }
 
 // BuildHandler is a Go build server LSP/JSON-RPC handler that wraps a
@@ -77,6 +80,7 @@ type BuildHandler struct {
 	init           *lspext.InitializeParams // set by "initialize" request
 	rootImportPath string                   // root import path of the workspace (e.g., "github.com/foo/bar")
 	cachingClient  *http.Client             // http.Client with a cache backed by an in-memory LRU cache
+	closers        []io.Closer              // values to dispose of when Close() is called
 }
 
 // reset clears all internal state in h.
@@ -102,7 +106,7 @@ func (h *BuildHandler) reset(init *lspext.InitializeParams, conn *jsonrpc2.Conn,
 	return nil
 }
 
-func (h *BuildHandler) handle(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) (result interface{}, err error) {
+func (h *BuildHandler) Handle(ctx context.Context, conn *jsonrpc2.Conn, req *jsonrpc2.Request) (result interface{}, err error) {
 	// Prevent any uncaught panics from taking the entire server down.
 	defer func() {
 		if r := recover(); r != nil {
@@ -187,10 +191,11 @@ func (h *BuildHandler) handle(ctx context.Context, conn *jsonrpc2.Conn, req *jso
 
 		// Determine the root import path of this workspace (e.g., "github.com/user/repo").
 		span.SetTag("originalRootPath", params.OriginalRootURI)
-		fs, err := RemoteFS(ctx, params)
+		fs, closer, err := RemoteFS(ctx, params)
 		if err != nil {
 			return nil, err
 		}
+		h.closers = append(h.closers, closer)
 
 		langInitParams, err := determineEnvironment(ctx, fs, params)
 		if err != nil {
@@ -571,3 +576,15 @@ func (h *BuildHandler) callLangServer(ctx context.Context, conn *jsonrpc2.Conn,
 	}
 	return nil
 }
+
+// Close implements io.Closer
+func (h *BuildHandler) Close() error {
+	var result error
+	for _, closer := range h.closers {
+		err := closer.Close()
+		if err != nil {
+			result = multierror.Append(result, err)
+		}
+	}
+	return result
+}

buildserver/build_server.go

@@ -3,6 +3,8 @@ package buildserver_test
 import (
 	"context"
 	"fmt"
+	"io"
+	"io/ioutil"
 	"net/url"
 	"os"
 	"strings"
@@ -267,15 +269,16 @@ func TestIntegration(t *testing.T) {
 // test.
 func useGithubForVFS() func() {
 	origRemoteFS := gobuildserver.RemoteFS
-	gobuildserver.RemoteFS = func(ctx context.Context, initializeParams lspext.InitializeParams) (ctxvfs.FileSystem, error) {
+	gobuildserver.RemoteFS = func(ctx context.Context, initializeParams lspext.InitializeParams) (ctxvfs.FileSystem, io.Closer, error) {
 		u, err := gituri.Parse(string(initializeParams.OriginalRootURI))
 		if err != nil {
-			return nil, errors.Wrap(err, "could not parse workspace URI for remotefs")
+			return nil, ioutil.NopCloser(strings.NewReader("")), errors.Wrap(err, "could not parse workspace URI for remotefs")
 		}
 		if u.Rev() == "" {
-			return nil, errors.Errorf("rev is required in uri: %s", initializeParams.OriginalRootURI)
+			return nil, ioutil.NopCloser(strings.NewReader("")), errors.Errorf("rev is required in uri: %s", initializeParams.OriginalRootURI)
 		}
-		return vfsutil.NewGitHubRepoVFS(ctx, string(u.Repo()), u.Rev())
+		fs, err := vfsutil.NewGitHubRepoVFS(ctx, string(u.Repo()), u.Rev())
+		return fs, ioutil.NopCloser(strings.NewReader("")), err
 	}
 
 	return func() {

buildserver/integration_test.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"net/url"
 	"os"
 	"path"
@@ -586,8 +587,8 @@ func yza() {}
 				}()
 
 				origRemoteFS := gobuildserver.RemoteFS
-				gobuildserver.RemoteFS = func(ctx context.Context, initializeParams lspext.InitializeParams) (ctxvfs.FileSystem, error) {
-					return mapFS(test.fs), nil
+				gobuildserver.RemoteFS = func(ctx context.Context, initializeParams lspext.InitializeParams) (ctxvfs.FileSystem, io.Closer, error) {
+					return mapFS(test.fs), ioutil.NopCloser(strings.NewReader("")), nil
 				}
 
 				defer func() {
@@ -703,7 +704,7 @@ func connectionToNewBuildServer(root string, t testing.TB) (*jsonrpc2.Conn, func
 	// do not use the pkg cache because tests won't install any pkgs
 	config.UseBinaryPkgCache = false
 
-	jsonrpc2.NewConn(context.Background(), a, jsonrpc2.AsyncHandler(gobuildserver.NewHandler(config)))
+	jsonrpc2.NewConn(context.Background(), a, jsonrpc2.AsyncHandler(jsonrpc2.HandlerWithError(gobuildserver.NewHandler(config).Handle)))
 
 	conn := jsonrpc2.NewConn(context.Background(), b, NoopHandler{}, jsonrpc2.OnRecv(onRecv), jsonrpc2.OnSend(onSend))
 	done := func() {

buildserver/proxy_test.go

@@ -13,11 +13,11 @@ import (
 
 	"github.com/neelance/parallel"
 	"github.com/sourcegraph/ctxvfs"
+	gobuildserver "github.com/sourcegraph/go-langserver/buildserver"
+	"github.com/sourcegraph/go-langserver/gituri"
 	"github.com/sourcegraph/go-langserver/pkg/lsp"
 	"github.com/sourcegraph/go-lsp/lspext"
 	"github.com/sourcegraph/jsonrpc2"
-	gobuildserver "github.com/sourcegraph/go-langserver/buildserver"
-	"github.com/sourcegraph/go-langserver/gituri"
 )
 
 // BenchmarkStress benchmarks performing "textDocument/definition",
@@ -60,7 +60,7 @@ func BenchmarkStress(b *testing.B) {
 		label := strings.Replace(root.Host+root.Path, "/", "-", -1)
 
 		b.Run(label, func(b *testing.B) {
-			fs, err := gobuildserver.RemoteFS(context.Background(), lspext.InitializeParams{OriginalRootURI: rootURI})
+			fs, _, err := gobuildserver.RemoteFS(context.Background(), lspext.InitializeParams{OriginalRootURI: rootURI})
 			if err != nil {
 				b.Fatal(err)
 			}

buildserver/stress_test.go

@@ -2,6 +2,9 @@ package buildserver
 
 import (
 	"context"
+	"io"
+	"io/ioutil"
+	"strings"
 
 	"github.com/pkg/errors"
 	"github.com/prometheus/client_golang/prometheus"
@@ -17,7 +20,7 @@ import (
 // SECURITY NOTE: This DOES NOT check that the user or context has permissions
 // to read the repo. We assume permission checks happen before a request reaches
 // a build server.
-var RemoteFS = func(ctx context.Context, initializeParams lspext.InitializeParams) (ctxvfs.FileSystem, error) {
+var RemoteFS = func(ctx context.Context, initializeParams lspext.InitializeParams) (ctxvfs.FileSystem, io.Closer, error) {
 	zipURL := func() string {
 		initializationOptions, ok := initializeParams.InitializationOptions.(map[string]interface{})
 		if !ok {
@@ -27,9 +30,10 @@ var RemoteFS = func(ctx context.Context, initializeParams lspext.InitializeParam
 		return url
 	}()
 	if zipURL != "" {
-		return vfsutil.NewZipVFS(ctx, zipURL, zipFetch.Inc, zipFetchFailed.Inc, true)
+		vfs, err := vfsutil.NewZipVFS(ctx, zipURL, zipFetch.Inc, zipFetchFailed.Inc, true)
+		return vfs, vfs, err
 	}
-	return nil, errors.Errorf("no zipURL was provided in the initializationOptions")
+	return nil, ioutil.NopCloser(strings.NewReader("")), errors.Errorf("no zipURL was provided in the initializationOptions")
 }
 
 var zipFetch = prometheus.NewCounter(prometheus.CounterOpts{

buildserver/vfs.go

@@ -35,6 +35,10 @@ type Store struct {
 	// BeforeEvict, when non-nil, is a function to call before evicting a file.
 	// It is passed the path to the file to be evicted.
 	BeforeEvict func(string)
+
+	// MaxCacheSizeBytes is the maximum size of the cache directory after evicting
+	// entries.
+	MaxCacheSizeBytes int64
 }
 
 // File is an os.File, but includes the Path
@@ -184,9 +188,15 @@ type EvictStats struct {
 	Evicted int
 }
 
+// Evict will remove files from Store.Dir until it is smaller than
+// Store.MaxCacheSizeBytes. It evicts files with the oldest modification time first.
+func (s *Store) Evict() {
+	s.EvictMaxSize(s.MaxCacheSizeBytes)
+}
+
 // Evict will remove files from Store.Dir until it is smaller than
 // maxCacheSizeBytes. It evicts files with the oldest modification time first.
-func (s *Store) Evict(maxCacheSizeBytes int64) (stats EvictStats, err error) {
+func (s *Store) EvictMaxSize(maxCacheSizeBytes int64) (stats EvictStats, err error) {
 	isZip := func(fi os.FileInfo) bool {
 		return strings.HasSuffix(fi.Name(), ".zip")
 	}

diskcache/cache.go

@@ -6,12 +6,14 @@ import (
 	"flag"
 	"fmt"
 	"io"
+	"io/ioutil"
 	"log"
 	"net"
 	"net/http"
 	"os"
 	"path/filepath"
 	"runtime/debug"
+	"strings"
 	"time"
 
 	"github.com/keegancsmith/tmpfriend"
@@ -33,15 +35,16 @@ import (
 )
 
 var (
-	mode           = flag.String("mode", "stdio", "communication mode (stdio|tcp|websocket)")
-	addr           = flag.String("addr", ":4389", "server listen address (tcp or websocket)")
-	trace          = flag.Bool("trace", false, "print all requests and responses")
-	logfile        = flag.String("logfile", "", "also log to this file (in addition to stderr)")
-	printVersion   = flag.Bool("version", false, "print version and exit")
-	pprof          = flag.String("pprof", "", "start a pprof http server (https://golang.org/pkg/net/http/pprof/)")
-	freeosmemory   = flag.Bool("freeosmemory", true, "aggressively free memory back to the OS")
-	useBuildServer = flag.Bool("usebuildserver", false, "use a build server to fetch dependencies, fetch files via Zip URL, etc.")
-	cacheDir       = flag.String("cachedir", "/tmp", "directory to store cached archives")
+	mode              = flag.String("mode", "stdio", "communication mode (stdio|tcp|websocket)")
+	addr              = flag.String("addr", ":4389", "server listen address (tcp or websocket)")
+	trace             = flag.Bool("trace", false, "print all requests and responses")
+	logfile           = flag.String("logfile", "", "also log to this file (in addition to stderr)")
+	printVersion      = flag.Bool("version", false, "print version and exit")
+	pprof             = flag.String("pprof", "", "start a pprof http server (https://golang.org/pkg/net/http/pprof/)")
+	freeosmemory      = flag.Bool("freeosmemory", true, "aggressively free memory back to the OS")
+	useBuildServer    = flag.Bool("usebuildserver", false, "use a build server to fetch dependencies, fetch files via Zip URL, etc.")
+	cacheDir          = flag.String("cachedir", "/tmp", "directory to store cached archives")
+	maxCacheSizeBytes = flag.Int64("maxCacheSizeBytes", 50*1024*1024*1024, "the maximum size of the cache directory after evicting entries")
 
 	// Default Config, can be overridden by InitializationOptions
 	usebinarypkgcache  = flag.Bool("usebinarypkgcache", true, "use $GOPATH/pkg binary .a files (improves performance). Can be overridden by InitializationOptions.")
@@ -76,6 +79,7 @@ func main() {
 	log.SetFlags(0)
 
 	vfsutil.ArchiveCacheDir = filepath.Join(*cacheDir, "lang-go-archive-cache")
+	vfsutil.MaxCacheSizeBytes = *maxCacheSizeBytes
 
 	// Start pprof server, if desired.
 	if *pprof != "" {
@@ -166,11 +170,12 @@ func run(cfg langserver.Config) error {
 		connOpt = append(connOpt, jsonrpc2.LogMessages(log.New(logW, "", 0)))
 	}
 
-	newHandler := func() jsonrpc2.Handler {
+	newHandler := func() (jsonrpc2.Handler, io.Closer) {
 		if *useBuildServer {
-			return jsonrpc2.AsyncHandler(buildserver.NewHandler(cfg))
+			handler := buildserver.NewHandler(cfg)
+			return jsonrpc2.AsyncHandler(jsonrpc2.HandlerWithError(handler.Handle)), handler
 		}
-		return langserver.NewHandler(cfg)
+		return langserver.NewHandler(cfg), ioutil.NopCloser(strings.NewReader(""))
 	}
 
 	switch *mode {
@@ -194,9 +199,14 @@ func run(cfg langserver.Config) error {
 			connectionID := connectionCount
 			log.Printf("langserver-go: received incoming connection #%d\n", connectionID)
 			openGauge.Inc()
-			jsonrpc2Connection := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewBufferedStream(conn, jsonrpc2.VSCodeObjectCodec{}), newHandler(), connOpt...)
+			handler, closer := newHandler()
+			jsonrpc2Connection := jsonrpc2.NewConn(context.Background(), jsonrpc2.NewBufferedStream(conn, jsonrpc2.VSCodeObjectCodec{}), handler, connOpt...)
 			go func() {
 				<-jsonrpc2Connection.DisconnectNotify()
+				err := closer.Close()
+				if err != nil {
+					log.Println(err)
+				}
 				log.Printf("langserver-go: connection #%d closed\n", connectionID)
 				openGauge.Dec()
 			}()
@@ -221,7 +231,12 @@ func run(cfg langserver.Config) error {
 
 			openGauge.Inc()
 			log.Printf("langserver-go: received incoming connection #%d\n", connectionID)
-			<-jsonrpc2.NewConn(context.Background(), wsjsonrpc2.NewObjectStream(connection), newHandler(), connOpt...).DisconnectNotify()
+			handler, closer := newHandler()
+			<-jsonrpc2.NewConn(context.Background(), wsjsonrpc2.NewObjectStream(connection), handler, connOpt...).DisconnectNotify()
+			err = closer.Close()
+			if err != nil {
+				log.Println(err)
+			}
 			log.Printf("langserver-go: connection #%d closed\n", connectionID)
 			openGauge.Dec()
 		})
@@ -243,7 +258,12 @@ func run(cfg langserver.Config) error {
 
 	case "stdio":
 		log.Println("langserver-go: reading on stdin, writing on stdout")
-		<-jsonrpc2.NewConn(context.Background(), jsonrpc2.NewBufferedStream(stdrwc{}, jsonrpc2.VSCodeObjectCodec{}), newHandler(), connOpt...).DisconnectNotify()
+		handler, closer := newHandler()
+		<-jsonrpc2.NewConn(context.Background(), jsonrpc2.NewBufferedStream(stdrwc{}, jsonrpc2.VSCodeObjectCodec{}), handler, connOpt...).DisconnectNotify()
+		err := closer.Close()
+		if err != nil {
+			log.Println(err)
+		}
 		log.Println("connection closed")
 		return nil