Skip to content

Commit

Permalink
Introduce RPCContext to aid more fine-grained & efficient processing (
Browse files Browse the repository at this point in the history 
#1402)

* Introduce RPCContext to aid more fine-grained & efficient processing

* Add RPCContext to job-related tests (#1371)

---------

Co-authored-by: Daniel Banck <[email protected]>
  • Loading branch information
@radeksimko @dbanck
radeksimko and dbanck authored Sep 13, 2023
1 parent b8df81a commit e22a021
Show file tree
Hide file tree
Showing 8 changed files with 77 additions and 7 deletions.
21 changes: 21 additions & 0 deletions internal/context/context.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -20,6 +20,18 @@ func (k *contextKey) String() string {
return k.Name
}

type RPCContextData struct {
Method string
URI string
}

func (rpcc RPCContextData) Copy() RPCContextData {
return RPCContextData{
Method: rpcc.Method,
URI: rpcc.URI,
}
}

var (
ctxTfExecPath = &contextKey{"terraform executable path"}
ctxTfExecLogPath = &contextKey{"terraform executor log path"}
Expand All @@ -30,6 +42,7 @@ var (
ctxLsVersion = &contextKey{"language server version"}
ctxProgressToken = &contextKey{"progress token"}
ctxExperimentalFeatures = &contextKey{"experimental features"}
ctxRPCContext = &contextKey{"rpc context"}
)

func missingContextErr(ctxKey *contextKey) *MissingContextErr {
Expand Down Expand Up @@ -165,3 +178,11 @@ 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 RPCContext(ctx context.Context) RPCContextData {
return ctx.Value(ctxRPCContext).(RPCContextData)
}
2 changes: 2 additions & 0 deletions internal/langserver/handlers/initialize.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -164,6 +164,8 @@ 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))

err = svc.closedDirWalker.StartWalking(walkerCtx)
if err != nil {
return serverCaps, fmt.Errorf("failed to start closedDirWalker: %w", err)
Expand Down
29 changes: 22 additions & 7 deletions internal/langserver/handlers/service.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -606,24 +606,39 @@ 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"`
}
type p struct {
URI string `json:"uri,omitempty"`
RootURI string `json:"rootUri,omitempty"`
TextDocument t `json:"textDocument,omitempty"`
RootURI string `json:"rootUri,omitempty"`
}
params := p{}
err := req.UnmarshalParams(&params)
if err == nil {
attrs = append(attrs, attribute.KeyValue{
Key: attribute.Key("URI"),
Value: attribute.StringValue(string(params.URI)),
})
if err != nil {
return nil, err
}

uri := params.TextDocument.URI
if params.RootURI != "" {
uri = params.RootURI
}

attrs = append(attrs, attribute.KeyValue{
Key: attribute.Key("URI"),
Value: attribute.StringValue(uri),
})

tracer := otel.Tracer(tracerName)
ctx, span := tracer.Start(ctx, "rpc:"+req.Method(),
trace.WithAttributes(attrs...))
defer span.End()

ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{
Method: req.Method(),
URI: uri,
})

result, err := rpch.New(fn)(ctx, req)
if ctx.Err() != nil && errors.Is(ctx.Err(), context.Canceled) {
err = fmt.Errorf("%w: %s", requestCancelled.Err(), err)
Expand Down
8 changes: 8 additions & 0 deletions internal/scheduler/scheduler_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import (
"time"

"github.com/google/go-cmp/cmp"
lsctx "github.com/hashicorp/terraform-ls/internal/context"
"github.com/hashicorp/terraform-ls/internal/document"
"github.com/hashicorp/terraform-ls/internal/job"
"github.com/hashicorp/terraform-ls/internal/state"
Expand All @@ -30,6 +31,7 @@ func TestScheduler_withIgnoreExistingState(t *testing.T) {
tmpDir := t.TempDir()

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})

s := NewScheduler(ss.JobStore, 1, job.LowPriority)
s.SetLogger(testLogger())
Expand Down Expand Up @@ -92,6 +94,7 @@ func TestScheduler_closedOnly(t *testing.T) {
tmpDir := t.TempDir()

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})

s := NewScheduler(ss.JobStore, 2, job.LowPriority)
s.SetLogger(testLogger())
Expand Down Expand Up @@ -154,6 +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{})
newId, err := ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(c context.Context) error {
atomic.AddInt64(&closedJobsExecuted, 1)
Expand All @@ -180,6 +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{})
newId, err := ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(c context.Context) error {
atomic.AddInt64(&openJobsExecuted, 1)
Expand Down Expand Up @@ -264,6 +269,7 @@ func BenchmarkScheduler_EnqueueAndWaitForJob_closedOnly(b *testing.B) {

tmpDir := b.TempDir()
ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})

s := NewScheduler(ss.JobStore, 1, job.LowPriority)
s.Start(ctx)
Expand Down Expand Up @@ -305,6 +311,7 @@ func TestScheduler_defer(t *testing.T) {
tmpDir := t.TempDir()

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})

s := NewScheduler(ss.JobStore, 2, job.LowPriority)
s.SetLogger(testLogger())
Expand Down Expand Up @@ -394,6 +401,7 @@ func TestScheduler_dependsOn(t *testing.T) {
tmpDir := t.TempDir()

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})

s := NewScheduler(ss.JobStore, 2, job.LowPriority)
s.SetLogger(testLogger())
Expand Down
6 changes: 6 additions & 0 deletions internal/state/jobs.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -13,6 +13,7 @@ import (
"time"

"github.com/hashicorp/go-memdb"
lsctx "github.com/hashicorp/terraform-ls/internal/context"
"github.com/hashicorp/terraform-ls/internal/document"
"github.com/hashicorp/terraform-ls/internal/job"
"go.opentelemetry.io/otel"
Expand Down Expand Up @@ -49,6 +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
}

func (sj *ScheduledJob) Copy() *ScheduledJob {
Expand All @@ -66,6 +69,7 @@ func (sj *ScheduledJob) Copy() *ScheduledJob {
DeferredJobIDs: sj.DeferredJobIDs.Copy(),
EnqueueTime: sj.EnqueueTime,
TraceSpan: traceSpan,
RPCContext: sj.RPCContext.Copy(),
}
}

Expand Down Expand Up @@ -125,6 +129,7 @@ func (js *JobStore) EnqueueJob(ctx context.Context, newJob job.Job) (job.ID, err
State: StateQueued,
EnqueueTime: time.Now(),
TraceSpan: jobSpan,
RPCContext: lsctx.RPCContext(ctx),
}

err := txn.Insert(js.tableName, sJob)
Expand Down Expand Up @@ -316,6 +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 = trace.ContextWithSpan(ctx, sJob.TraceSpan)

_, span := otel.Tracer(tracerName).Start(ctx, "job-wait",
Expand Down
13 changes: 13 additions & 0 deletions internal/state/jobs_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@ import (
"time"

"github.com/google/go-cmp/cmp"
lsctx "github.com/hashicorp/terraform-ls/internal/context"
"github.com/hashicorp/terraform-ls/internal/document"
"github.com/hashicorp/terraform-ls/internal/job"
)
Expand All @@ -27,6 +28,7 @@ func TestJobStore_EnqueueJob(t *testing.T) {
}

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
return nil
Expand Down Expand Up @@ -81,6 +83,7 @@ func TestJobStore_EnqueueJob_openDir(t *testing.T) {
}

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
id, err := ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
return nil
Expand Down Expand Up @@ -159,6 +162,7 @@ func TestJobStore_EnqueueJob_verify(t *testing.T) {

jobCount := 50
ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})

for i := 0; i < jobCount; i++ {
i := i
Expand Down Expand Up @@ -217,6 +221,7 @@ func TestJobStore_DequeueJobsForDir(t *testing.T) {
}

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
firstDir := document.DirHandleFromPath("/test-1")
_, err = ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
Expand Down Expand Up @@ -261,6 +266,7 @@ func TestJobStore_AwaitNextJob_closedOnly(t *testing.T) {
}

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
firstDir := document.DirHandleFromPath("/test-1")
id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
Expand Down Expand Up @@ -324,6 +330,7 @@ func TestJobStore_AwaitNextJob_openOnly(t *testing.T) {
}

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
firstDir := document.DirHandleFromPath("/test-1")
_, err = ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
Expand Down Expand Up @@ -387,6 +394,7 @@ func TestJobStore_AwaitNextJob_highPriority(t *testing.T) {
}

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
firstDir := document.DirHandleFromPath("/test-1")
id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
Expand Down Expand Up @@ -468,6 +476,7 @@ func TestJobStore_AwaitNextJob_lowPriority(t *testing.T) {
}

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
firstDir := document.DirHandleFromPath("/test-1")
id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
Expand Down Expand Up @@ -564,6 +573,7 @@ func TestJobStore_WaitForJobs(t *testing.T) {
}

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
return nil
Expand Down Expand Up @@ -605,6 +615,7 @@ func TestJobStore_FinishJob_basic(t *testing.T) {
}

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
return nil
Expand Down Expand Up @@ -666,6 +677,7 @@ func TestJobStore_FinishJob_defer(t *testing.T) {
}

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
id1, err := ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
return nil
Expand Down Expand Up @@ -719,6 +731,7 @@ func TestJobStore_FinishJob_dependsOn(t *testing.T) {
}

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
parentId, err := ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
return nil
Expand Down
2 changes: 2 additions & 0 deletions internal/state/module_changes_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,7 @@ import (

"github.com/google/go-cmp/cmp"
"github.com/hashicorp/go-version"
lsctx "github.com/hashicorp/terraform-ls/internal/context"
"github.com/hashicorp/terraform-ls/internal/document"
"github.com/hashicorp/terraform-ls/internal/job"
tfaddr "github.com/hashicorp/terraform-registry-address"
Expand Down Expand Up @@ -96,6 +97,7 @@ func TestModuleChanges_AwaitNextChangeBatch_maxTimespan(t *testing.T) {
modHandle := document.DirHandleFromPath(modPath)

ctx := context.Background()
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
_, err = ss.JobStore.EnqueueJob(ctx, job.Job{
Func: func(ctx context.Context) error {
return nil
Expand Down
3 changes: 3 additions & 0 deletions internal/walker/walker_test.go
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,6 +15,7 @@ import (
"github.com/google/go-cmp/cmp"
"github.com/hashicorp/go-version"
tfjson "github.com/hashicorp/terraform-json"
lsctx "github.com/hashicorp/terraform-ls/internal/context"
"github.com/hashicorp/terraform-ls/internal/document"
"github.com/hashicorp/terraform-ls/internal/filesystem"
"github.com/hashicorp/terraform-ls/internal/indexer"
Expand Down Expand Up @@ -55,6 +56,7 @@ func TestWalker_basic(t *testing.T) {
t.Fatal(err)
}

ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
err = w.StartWalking(ctx)
if err != nil {
t.Fatal(err)
Expand Down Expand Up @@ -399,6 +401,7 @@ func TestWalker_complexModules(t *testing.T) {
if err != nil {
t.Fatal(err)
}
ctx = lsctx.WithRPCContext(ctx, lsctx.RPCContextData{})
err = w.StartWalking(ctx)
if err != nil {
t.Fatal(err)
Expand Down

0 comments on commit e22a021

Please sign in to comment.