context: Replace RPCContextData with DocumentContext

internal/context/context.go

@@ -20,15 +20,17 @@ func (k *contextKey) String() string {
 	return k.Name
 }
 
-type RPCContextData struct {
-	Method string
-	URI    string
+type Document struct {
+	Method     string
+	LanguageID string
+	URI        string
 }
 
-func (rpcc RPCContextData) Copy() RPCContextData {
-	return RPCContextData{
-		Method: rpcc.Method,
-		URI:    rpcc.URI,
+func (rpcc Document) Copy() Document {
+	return Document{
+		Method:     rpcc.Method,
+		LanguageID: rpcc.LanguageID,
+		URI:        rpcc.URI,
 	}
 }
 
@@ -42,8 +44,7 @@ var (
 	ctxLsVersion            = &contextKey{"language server version"}
 	ctxProgressToken        = &contextKey{"progress token"}
 	ctxExperimentalFeatures = &contextKey{"experimental features"}
-	ctxRPCContext           = &contextKey{"rpc context"}
-	ctxLanguageId           = &contextKey{"language ID"}
+	ctxDocumentContext      = &contextKey{"rpc context"}
 	ctxValidationOptions    = &contextKey{"validation options"}
 )
 
@@ -181,30 +182,18 @@ func ExperimentalFeatures(ctx context.Context) (settings.ExperimentalFeatures, e
 	return *expFeatures, nil
 }
 
-func WithRPCContext(ctx context.Context, rpcc RPCContextData) context.Context {
-	return context.WithValue(ctx, ctxRPCContext, rpcc)
+func WithDocumentContext(ctx context.Context, rpcc Document) context.Context {
+	return context.WithValue(ctx, ctxDocumentContext, rpcc)
 }
 
-func RPCContext(ctx context.Context) RPCContextData {
-	return ctx.Value(ctxRPCContext).(RPCContextData)
+func DocumentContext(ctx context.Context) Document {
+	return ctx.Value(ctxDocumentContext).(Document)
 }
 
-func (ctxData RPCContextData) IsDidChangeRequest() bool {
+func (ctxData Document) IsDidChangeRequest() bool {
 	return ctxData.Method == "textDocument/didChange"
 }
 
-func WithLanguageId(ctx context.Context, languageId string) context.Context {
-	return context.WithValue(ctx, ctxLanguageId, languageId)
-}
-
-func IsLanguageId(ctx context.Context, expectedLangId string) bool {
-	langId, ok := ctx.Value(ctxLanguageId).(string)
-	if !ok {
-		return false
-	}
-	return langId == expectedLangId
-}
-
 func WithValidationOptions(ctx context.Context, validationOptions *settings.ValidationOptions) context.Context {
 	return context.WithValue(ctx, ctxValidationOptions, validationOptions)
 }

internal/decoder/decoder_test.go

@@ -61,7 +61,7 @@ func TestDecoder_CodeLensesForFile_concurrencyBug(t *testing.T) {
 		if err != nil {
 			t.Error(err)
 		}
-		ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+		ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 		err = module.ParseModuleConfiguration(ctx, mapFs, ss.Modules, dirName)
 		if err != nil {
 			t.Error(err)

internal/langserver/handlers/did_change.go

@@ -29,7 +29,9 @@ func (svc *service) TextDocumentDidChange(ctx context.Context, params lsp.DidCha
 		return err
 	}
 
-	ctx = lsctx.WithLanguageId(ctx, doc.LanguageID)
+	docCtx := lsctx.DocumentContext(ctx)
+	docCtx.LanguageID = doc.LanguageID
+	ctx = lsctx.WithDocumentContext(ctx, docCtx)
 
 	newVersion := int(p.TextDocument.Version)
 

internal/langserver/handlers/did_open.go

@@ -8,7 +8,6 @@ import (
 	"fmt"
 
 	"github.com/creachadair/jrpc2"
-	lsctx "github.com/hashicorp/terraform-ls/internal/context"
 	"github.com/hashicorp/terraform-ls/internal/document"
 	lsp "github.com/hashicorp/terraform-ls/internal/protocol"
 	"github.com/hashicorp/terraform-ls/internal/state"
@@ -38,8 +37,6 @@ func (svc *service) TextDocumentDidOpen(ctx context.Context, params lsp.DidOpenT
 		return err
 	}
 
-	ctx = lsctx.WithLanguageId(ctx, params.TextDocument.LanguageID)
-
 	mod, err := svc.modStore.ModuleByPath(dh.Dir.Path())
 	if err != nil {
 		if state.IsModuleNotFound(err) {

internal/langserver/handlers/initialize.go

@@ -166,7 +166,7 @@ func (svc *service) Initialize(ctx context.Context, params lsp.InitializeParams)
 	// passing the request context here
 	// Static user-provided paths take precedence over dynamic discovery
 	walkerCtx := context.Background()
-	walkerCtx = lsctx.WithRPCContext(walkerCtx, lsctx.RPCContext(ctx))
+	walkerCtx = lsctx.WithDocumentContext(walkerCtx, lsctx.DocumentContext(ctx))
 
 	err = svc.closedDirWalker.StartWalking(walkerCtx)
 	if err != nil {

internal/langserver/handlers/service.go

@@ -614,7 +614,8 @@ func handle(ctx context.Context, req *jrpc2.Request, fn interface{}) (interface{
 	// We could capture all parameters here but for now we just
 	// opportunistically track the most important ones only.
 	type t struct {
-		URI string `json:"uri,omitempty"`
+		URI        string `json:"uri,omitempty"`
+		LanguageID string `json:"languageId,omitempty"`
 	}
 	type p struct {
 		TextDocument t      `json:"textDocument,omitempty"`
@@ -641,9 +642,10 @@ func handle(ctx context.Context, req *jrpc2.Request, fn interface{}) (interface{
 		trace.WithAttributes(attrs...))
 	defer span.End()
 
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{
-		Method: req.Method(),
-		URI:    uri,
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{
+		Method:     req.Method(),
+		LanguageID: params.TextDocument.LanguageID,
+		URI:        uri,
 	})
 
 	result, err := rpch.New(fn)(ctx, req)

internal/scheduler/scheduler_test.go

@@ -31,7 +31,7 @@ func TestScheduler_withIgnoreExistingState(t *testing.T) {
 	tmpDir := t.TempDir()
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 
 	s := NewScheduler(ss.JobStore, 1, job.LowPriority)
 	s.SetLogger(testLogger())
@@ -94,7 +94,7 @@ func TestScheduler_closedOnly(t *testing.T) {
 	tmpDir := t.TempDir()
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 
 	s := NewScheduler(ss.JobStore, 2, job.LowPriority)
 	s.SetLogger(testLogger())
@@ -157,7 +157,7 @@ func TestScheduler_closedAndOpen(t *testing.T) {
 			dirPath := filepath.Join(tmpDir, fmt.Sprintf("folder-x-%d", i))
 
 			ctx := context.Background()
-			ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+			ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 			newId, err := ss.JobStore.EnqueueJob(ctx, job.Job{
 				Func: func(c context.Context) error {
 					atomic.AddInt64(&closedJobsExecuted, 1)
@@ -184,7 +184,7 @@ func TestScheduler_closedAndOpen(t *testing.T) {
 			dirPath := filepath.Join(tmpDir, fmt.Sprintf("folder-y-%d", i))
 
 			ctx := context.Background()
-			ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+			ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 			newId, err := ss.JobStore.EnqueueJob(ctx, job.Job{
 				Func: func(c context.Context) error {
 					atomic.AddInt64(&openJobsExecuted, 1)
@@ -269,7 +269,7 @@ func BenchmarkScheduler_EnqueueAndWaitForJob_closedOnly(b *testing.B) {
 
 	tmpDir := b.TempDir()
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 
 	s := NewScheduler(ss.JobStore, 1, job.LowPriority)
 	s.Start(ctx)
@@ -311,7 +311,7 @@ func TestScheduler_defer(t *testing.T) {
 	tmpDir := t.TempDir()
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 
 	s := NewScheduler(ss.JobStore, 2, job.LowPriority)
 	s.SetLogger(testLogger())
@@ -401,7 +401,7 @@ func TestScheduler_dependsOn(t *testing.T) {
 	tmpDir := t.TempDir()
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 
 	s := NewScheduler(ss.JobStore, 2, job.LowPriority)
 	s.SetLogger(testLogger())

internal/state/jobs.go

@@ -50,8 +50,8 @@ type ScheduledJob struct {
 	// TraceSpan represents a tracing span for the entire job lifecycle
 	// (from queuing to finishing execution).
 	TraceSpan trace.Span
-	// RPCContext contains information from when & where the job was scheduled from
-	RPCContext lsctx.RPCContextData
+	// DocumentContext contains information from when & where the job was scheduled from
+	DocumentContext lsctx.Document
 }
 
 func (sj *ScheduledJob) Copy() *ScheduledJob {
@@ -61,15 +61,15 @@ func (sj *ScheduledJob) Copy() *ScheduledJob {
 	traceSpan := trace.SpanFromContext(newCtx)
 
 	return &ScheduledJob{
-		ID:             sj.ID,
-		Job:            sj.Job.Copy(),
-		IsDirOpen:      sj.IsDirOpen,
-		State:          sj.State,
-		JobErr:         sj.JobErr,
-		DeferredJobIDs: sj.DeferredJobIDs.Copy(),
-		EnqueueTime:    sj.EnqueueTime,
-		TraceSpan:      traceSpan,
-		RPCContext:     sj.RPCContext.Copy(),
+		ID:              sj.ID,
+		Job:             sj.Job.Copy(),
+		IsDirOpen:       sj.IsDirOpen,
+		State:           sj.State,
+		JobErr:          sj.JobErr,
+		DeferredJobIDs:  sj.DeferredJobIDs.Copy(),
+		EnqueueTime:     sj.EnqueueTime,
+		TraceSpan:       traceSpan,
+		DocumentContext: sj.DocumentContext.Copy(),
 	}
 }
 
@@ -123,13 +123,13 @@ func (js *JobStore) EnqueueJob(ctx context.Context, newJob job.Job) (job.ID, err
 		}))
 
 	sJob := &ScheduledJob{
-		ID:          newJobID,
-		Job:         newJob,
-		IsDirOpen:   dirOpen,
-		State:       StateQueued,
-		EnqueueTime: time.Now(),
-		TraceSpan:   jobSpan,
-		RPCContext:  lsctx.RPCContext(ctx),
+		ID:              newJobID,
+		Job:             newJob,
+		IsDirOpen:       dirOpen,
+		State:           StateQueued,
+		EnqueueTime:     time.Now(),
+		TraceSpan:       jobSpan,
+		DocumentContext: lsctx.DocumentContext(ctx),
 	}
 
 	err := txn.Insert(js.tableName, sJob)
@@ -321,7 +321,7 @@ func (js *JobStore) awaitNextJob(ctx context.Context, priority job.JobPriority)
 	js.logger.Printf("JOBS: Dispatching next job %q (scheduler prio: %d, job prio: %d, isDirOpen: %t): %q for %q",
 		sJob.ID, priority, sJob.Priority, sJob.IsDirOpen, sJob.Type, sJob.Dir)
 
-	ctx = lsctx.WithRPCContext(ctx, sJob.RPCContext)
+	ctx = lsctx.WithDocumentContext(ctx, sJob.DocumentContext)
 	ctx = trace.ContextWithSpan(ctx, sJob.TraceSpan)
 
 	_, span := otel.Tracer(tracerName).Start(ctx, "job-wait",

internal/state/jobs_test.go

@@ -28,7 +28,7 @@ func TestJobStore_EnqueueJob(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
 			return nil
@@ -83,7 +83,7 @@ func TestJobStore_EnqueueJob_openDir(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	id, err := ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
 			return nil
@@ -162,7 +162,7 @@ func TestJobStore_EnqueueJob_verify(t *testing.T) {
 
 	jobCount := 50
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 
 	for i := 0; i < jobCount; i++ {
 		i := i
@@ -221,7 +221,7 @@ func TestJobStore_DequeueJobsForDir(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	firstDir := document.DirHandleFromPath("/test-1")
 	_, err = ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
@@ -266,7 +266,7 @@ func TestJobStore_AwaitNextJob_closedOnly(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	firstDir := document.DirHandleFromPath("/test-1")
 	id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
@@ -330,7 +330,7 @@ func TestJobStore_AwaitNextJob_openOnly(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	firstDir := document.DirHandleFromPath("/test-1")
 	_, err = ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
@@ -394,7 +394,7 @@ func TestJobStore_AwaitNextJob_highPriority(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	firstDir := document.DirHandleFromPath("/test-1")
 	id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
@@ -476,7 +476,7 @@ func TestJobStore_AwaitNextJob_lowPriority(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	firstDir := document.DirHandleFromPath("/test-1")
 	id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
@@ -573,7 +573,7 @@ func TestJobStore_WaitForJobs(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
 			return nil
@@ -615,7 +615,7 @@ func TestJobStore_FinishJob_basic(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
 			return nil
@@ -677,7 +677,7 @@ func TestJobStore_FinishJob_defer(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
 			return nil
@@ -731,7 +731,7 @@ func TestJobStore_FinishJob_dependsOn(t *testing.T) {
 	}
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	parentId, err := ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
 			return nil

internal/state/module_changes_test.go

@@ -97,7 +97,7 @@ func TestModuleChanges_AwaitNextChangeBatch_maxTimespan(t *testing.T) {
 	modHandle := document.DirHandleFromPath(modPath)
 
 	ctx := context.Background()
-	ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
+	ctx = lsctx.WithDocumentContext(ctx, lsctx.Document{})
 	_, err = ss.JobStore.EnqueueJob(ctx, job.Job{
 		Func: func(ctx context.Context) error {
 			return nil