1// Copyright 2020 The Go Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE file.
4
5package lsprpc
6
7import (
8	"context"
9	"encoding/json"
10	"time"
11
12	"golang.org/x/tools/internal/jsonrpc2"
13	"golang.org/x/tools/internal/lsp/telemetry"
14	"golang.org/x/tools/internal/telemetry/trace"
15)
16
17type telemetryHandler struct{}
18
19func (h telemetryHandler) Deliver(ctx context.Context, r *jsonrpc2.Request, delivered bool) bool {
20	stats := h.getStats(ctx)
21	if stats != nil {
22		stats.delivering()
23	}
24	return false
25}
26
27func (h telemetryHandler) Cancel(ctx context.Context, conn *jsonrpc2.Conn, id jsonrpc2.ID, cancelled bool) bool {
28	return false
29}
30
31func (h telemetryHandler) Request(ctx context.Context, conn *jsonrpc2.Conn, direction jsonrpc2.Direction, r *jsonrpc2.WireRequest) context.Context {
32	if r.Method == "" {
33		panic("no method in rpc stats")
34	}
35	stats := &rpcStats{
36		method:    r.Method,
37		start:     time.Now(),
38		direction: direction,
39		payload:   r.Params,
40	}
41	ctx = context.WithValue(ctx, statsKey, stats)
42	mode := telemetry.Outbound
43	if direction == jsonrpc2.Receive {
44		mode = telemetry.Inbound
45	}
46	ctx, stats.close = trace.StartSpan(ctx, r.Method,
47		telemetry.Method.Of(r.Method),
48		telemetry.RPCDirection.Of(mode),
49		telemetry.RPCID.Of(r.ID),
50	)
51	telemetry.Started.Record(ctx, 1)
52	_, stats.delivering = trace.StartSpan(ctx, "queued")
53	return ctx
54}
55
56func (h telemetryHandler) Response(ctx context.Context, conn *jsonrpc2.Conn, direction jsonrpc2.Direction, r *jsonrpc2.WireResponse) context.Context {
57	return ctx
58}
59
60func (h telemetryHandler) Done(ctx context.Context, err error) {
61	stats := h.getStats(ctx)
62	if err != nil {
63		ctx = telemetry.StatusCode.With(ctx, "ERROR")
64	} else {
65		ctx = telemetry.StatusCode.With(ctx, "OK")
66	}
67	elapsedTime := time.Since(stats.start)
68	latencyMillis := float64(elapsedTime) / float64(time.Millisecond)
69	telemetry.Latency.Record(ctx, latencyMillis)
70	stats.close()
71}
72
73func (h telemetryHandler) Read(ctx context.Context, bytes int64) context.Context {
74	telemetry.SentBytes.Record(ctx, bytes)
75	return ctx
76}
77
78func (h telemetryHandler) Wrote(ctx context.Context, bytes int64) context.Context {
79	telemetry.ReceivedBytes.Record(ctx, bytes)
80	return ctx
81}
82
83const eol = "\r\n\r\n\r\n"
84
85func (h telemetryHandler) Error(ctx context.Context, err error) {
86}
87
88func (h telemetryHandler) getStats(ctx context.Context) *rpcStats {
89	stats, ok := ctx.Value(statsKey).(*rpcStats)
90	if !ok || stats == nil {
91		method, ok := ctx.Value(telemetry.Method).(string)
92		if !ok {
93			method = "???"
94		}
95		stats = &rpcStats{
96			method: method,
97			close:  func() {},
98		}
99	}
100	return stats
101}
102
103type rpcStats struct {
104	method     string
105	direction  jsonrpc2.Direction
106	id         *jsonrpc2.ID
107	payload    *json.RawMessage
108	start      time.Time
109	delivering func()
110	close      func()
111}
112
113type statsKeyType int
114
115const statsKey = statsKeyType(0)
116