pkg/lsp: refactor diagnostics

hide-org · Oct 6, 2024 · 5657a42 · 5657a42
1 parent 9d57ec0
commit 5657a42
Show file tree

Hide file tree

Showing 4 changed files with 101 additions and 38 deletions.
diff --git a/cmd/run.go b/cmd/run.go
@@ -98,9 +98,9 @@ var runCmd = &cobra.Command{
 
 		fileManager := files.NewFileManager(gitignore.NewMatcherFactory())
 		languageDetector := lsp.NewLanguageDetector()
-		diagnosticsStore := lsp.NewDiagnosticsStore()
+		diagnosticsService := lsp.NewDiagnosticsService()
 		clientPool := lsp.NewClientPool()
-		lspService := lsp.NewService(languageDetector, lsp.LspServerExecutables, diagnosticsStore, clientPool)
+		lspService := lsp.NewService(languageDetector, lsp.LspServerExecutables, diagnosticsService, clientPool)
 		projectManager := project.NewProjectManager(containerRunner, projectStore, projectsDir, fileManager, lspService, languageDetector, random.String)
 		validator := validator.New(validator.WithRequiredStructEnabled())
 

diff --git a/pkg/lsp/client.go b/pkg/lsp/client.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"encoding/json"
 
+	"github.com/rs/zerolog/log"
 	"github.com/sourcegraph/jsonrpc2"
 
 	protocol "github.com/tliron/glsp/protocol_3_16"
@@ -44,6 +45,7 @@ type ClientImpl struct {
 func NewClient(server Process, diagnosticsChannel chan protocol.PublishDiagnosticsParams) Client {
 	handler := &lspHandler{
 		diagnosticsHandler: func(params protocol.PublishDiagnosticsParams) {
+			log.Debug().Msgf("Handling diagnostics: %+v", params)
 			diagnosticsChannel <- params
 		},
 	}
@@ -83,6 +85,5 @@ func (c *ClientImpl) NotifyDidClose(ctx context.Context, params protocol.DidClos
 
 func (c *ClientImpl) Shutdown(ctx context.Context) error {
 	err := c.server.Stop()
-	close(c.diagnosticsChannel)
 	return err
 }
diff --git a/pkg/lsp/diagnostics.go b/pkg/lsp/diagnostics.go
@@ -1,20 +1,36 @@
 package lsp
 
 import (
+	"context"
 	"sync"
 
+	"github.com/rs/zerolog/log"
 	protocol "github.com/tliron/glsp/protocol_3_16"
 )
 
+// TODO: update docs
 // In memory store for diagnostics. Applies mutex locking for concurrent access.
-type DiagnosticsStore struct {
+type DiagnosticsService struct {
 	diagnostics map[ProjectId]map[protocol.DocumentUri][]protocol.Diagnostic
-	mu          sync.Mutex
+	listeners   map[ProjectId]map[LanguageId]*listenerInfo
+	mu          sync.RWMutex
 }
 
-func (d *DiagnosticsStore) Get(projectId ProjectId, uri protocol.DocumentUri) ([]protocol.Diagnostic, bool) {
-	d.mu.Lock()
-	defer d.mu.Unlock()
+type listenerInfo struct {
+	cancel  context.CancelFunc
+	channel chan protocol.PublishDiagnosticsParams
+}
+
+func NewDiagnosticsService() *DiagnosticsService {
+	return &DiagnosticsService{
+		diagnostics: make(map[ProjectId]map[protocol.DocumentUri][]protocol.Diagnostic),
+		listeners:   make(map[ProjectId]map[LanguageId]*listenerInfo),
+	}
+}
+
+func (d *DiagnosticsService) Get(projectId ProjectId, uri protocol.DocumentUri) ([]protocol.Diagnostic, bool) {
+	d.mu.RLock()
+	defer d.mu.RUnlock()
 
 	if diagnostics, ok := d.diagnostics[projectId]; ok {
 		if diagnostics, ok := diagnostics[uri]; ok {
@@ -25,7 +41,7 @@ func (d *DiagnosticsStore) Get(projectId ProjectId, uri protocol.DocumentUri) ([
 	return nil, false
 }
 
-func (d *DiagnosticsStore) Set(projectId ProjectId, uri protocol.DocumentUri, diagnostics []protocol.Diagnostic) {
+func (d *DiagnosticsService) set(projectId ProjectId, uri protocol.DocumentUri, diagnostics []protocol.Diagnostic) {
 	d.mu.Lock()
 	defer d.mu.Unlock()
 
@@ -36,15 +52,76 @@ func (d *DiagnosticsStore) Set(projectId ProjectId, uri protocol.DocumentUri, di
 	d.diagnostics[projectId][uri] = diagnostics
 }
 
-func (d *DiagnosticsStore) DeleteAllForProject(projectId ProjectId) {
+func (d *DiagnosticsService) StartListener(projectId ProjectId, languageId LanguageId, channel chan protocol.PublishDiagnosticsParams) {
 	d.mu.Lock()
 	defer d.mu.Unlock()
 
-	delete(d.diagnostics, projectId)
+	if _, exists := d.listeners[projectId]; !exists {
+		d.listeners[projectId] = make(map[LanguageId]*listenerInfo)
+	}
+
+	// TODO: should we cancel or return error?
+	// Cancel existing listener if any
+	if info, exists := d.listeners[projectId][languageId]; exists {
+		info.cancel()
+		close(info.channel)
+	}
+
+	ctx, cancel := context.WithCancel(context.Background())
+	d.listeners[projectId][languageId] = &listenerInfo{
+		cancel:  cancel,
+		channel: channel,
+	}
+
+	go d.listen(ctx, projectId, languageId, channel)
 }
 
-func NewDiagnosticsStore() *DiagnosticsStore {
-	return &DiagnosticsStore{
-		diagnostics: make(map[ProjectId]map[protocol.DocumentUri][]protocol.Diagnostic),
+func (d *DiagnosticsService) StopListener(projectId ProjectId, languageId LanguageId) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	if projectListeners, exists := d.listeners[projectId]; exists {
+		if info, exists := projectListeners[languageId]; exists {
+			info.cancel()
+			close(info.channel)
+			delete(projectListeners, languageId)
+		}
+		if len(projectListeners) == 0 {
+			delete(d.listeners, projectId)
+		}
 	}
 }
+
+func (d *DiagnosticsService) listen(ctx context.Context, projectId ProjectId, languageId LanguageId, channel chan protocol.PublishDiagnosticsParams) {
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case diagnostics, ok := <-channel:
+			if !ok {
+				return
+			}
+			d.set(projectId, diagnostics.URI, diagnostics.Diagnostics)
+			log.Debug().
+				Str("projectId", string(projectId)).
+				Str("languageId", string(languageId)).
+				Str("uri", string(diagnostics.URI)).
+				Msgf("Received diagnostics: %d", len(diagnostics.Diagnostics))
+		}
+	}
+}
+
+func (d *DiagnosticsService) DeleteAllForProject(projectId ProjectId) {
+	d.mu.Lock()
+	defer d.mu.Unlock()
+
+	delete(d.diagnostics, projectId)
+	if projectListeners, exists := d.listeners[projectId]; exists {
+		for _, info := range projectListeners {
+			info.cancel()
+			close(info.channel)
+		}
+		delete(d.listeners, projectId)
+	}
+}
+
diff --git a/pkg/lsp/service.go b/pkg/lsp/service.go
@@ -39,7 +39,7 @@ type Service interface {
 type ServiceImpl struct {
 	languageDetector     LanguageDetector
 	clientPool           ClientPool
-	diagnosticsStore     *DiagnosticsStore
+	diagnosticsService   *DiagnosticsService
 	lspServerExecutables map[LanguageId]Command
 }
 
@@ -115,9 +115,7 @@ func (s *ServiceImpl) StartServer(ctx context.Context, languageId LanguageId) er
 	}
 
 	s.clientPool.Set(projectId, languageId, client)
-
-	// TODO: kill this goroutine when the project is deleted
-	go s.listenForDiagnostics(projectId, diagnosticsChannel)
+	s.diagnosticsService.StartListener(projectId, languageId, diagnosticsChannel)
 	return nil
 }
 
@@ -141,6 +139,9 @@ func (s *ServiceImpl) StopServer(ctx context.Context, languageId LanguageId) err
 	}
 
 	s.clientPool.Delete(project.Id, languageId)
+	s.diagnosticsService.StopListener(project.Id, languageId)
+	// TODO: do we need this?
+	// s.diagnosticsManager.DeleteAllForLanguage(project.Id, languageId)
 
 	return nil
 }
@@ -280,7 +281,7 @@ func (s *ServiceImpl) GetDiagnostics(ctx context.Context, file model.File) ([]pr
 	}
 
 	uri := PathToURI(filepath.Join(project.Path, file.Path))
-	if diagnostics, ok := s.diagnosticsStore.Get(project.Id, uri); ok {
+	if diagnostics, ok := s.diagnosticsService.Get(project.Id, uri); ok {
 		return diagnostics, nil
 	}
 
@@ -300,7 +301,7 @@ func (s *ServiceImpl) CleanupProject(ctx context.Context, projectId ProjectId) e
 	}
 
 	s.clientPool.DeleteAllForProject(projectId)
-	s.diagnosticsStore.DeleteAllForProject(projectId)
+	s.diagnosticsService.DeleteAllForProject(projectId)
 	return nil
 }
 
@@ -336,31 +337,15 @@ func (s *ServiceImpl) getClients(ctx context.Context) []Client {
 	return clients
 }
 
-func (s *ServiceImpl) listenForDiagnostics(projectId ProjectId, channel chan protocol.PublishDiagnosticsParams) {
-	for {
-		select {
-		case diagnostics := <-channel:
-			log.Debug().Str("projectId", projectId).Str("uri", diagnostics.URI).Msg("Received diagnostics")
-			log.Debug().Str("projectId", projectId).Str("uri", diagnostics.URI).Msgf("Diagnostics: %+v", diagnostics.Diagnostics)
-
-			s.updateDiagnostics(projectId, diagnostics)
-		}
-	}
-}
-
-func (s *ServiceImpl) updateDiagnostics(projectId ProjectId, diagnostics protocol.PublishDiagnosticsParams) {
-	s.diagnosticsStore.Set(projectId, diagnostics.URI, diagnostics.Diagnostics)
-}
-
 func PathToURI(path string) protocol.DocumentUri {
 	return protocol.DocumentUri("file://" + path)
 }
 
-func NewService(languageDetector LanguageDetector, lspServerExecutables map[LanguageId]Command, diagnosticsStore *DiagnosticsStore, clientPool ClientPool) Service {
+func NewService(languageDetector LanguageDetector, lspServerExecutables map[LanguageId]Command, diagnosticsService *DiagnosticsService, clientPool ClientPool) Service {
 	return &ServiceImpl{
 		languageDetector:     languageDetector,
 		clientPool:           clientPool,
-		diagnosticsStore:     diagnosticsStore,
+		diagnosticsService:   diagnosticsService,
 		lspServerExecutables: lspServerExecutables,
 	}
 }