Skip to content

Commit 1d5eca4

Browse files
ammarioammario
committed
Extract endpoints into new /metrics section
1 parent 939a07d commit 1d5eca4

File tree

6 files changed

+247
-191
lines changed

6 files changed

+247
-191
lines changed

‎coderd/coderd.go

+6-1
Original file line numberDiff line numberDiff line change
@@ -335,6 +335,12 @@ func New(options *Options) *API {
335335
})
336336
})
337337
})
338+
r.Route("/metrics", func(r chi.Router) {
339+
r.Group(func(r chi.Router) {
340+
r.Use(httpmw.ExtractWorkspaceAgent(options.Database))
341+
r.Get("/report-agent-stats", api.workspaceAgentReportStats)
342+
})
343+
})
338344
r.Route("/workspaceagents", func(r chi.Router) {
339345
r.Post("/azure-instance-identity", api.postWorkspaceAuthAzureInstanceIdentity)
340346
r.Post("/aws-instance-identity", api.postWorkspaceAuthAWSInstanceIdentity)
@@ -343,7 +349,6 @@ func New(options *Options) *API {
343349
r.Use(httpmw.ExtractWorkspaceAgent(options.Database))
344350
r.Get("/metadata", api.workspaceAgentMetadata)
345351
r.Get("/listen", api.workspaceAgentListen)
346-
r.Get("/report-stats", api.workspaceAgentReportStats)
347352

348353
r.Get("/gitsshkey", api.agentGitSSHKey)
349354
r.Get("/turn", api.workspaceAgentTurn)

‎coderd/metrics.go

+146
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,146 @@
1+
package coderd
2+
3+
import (
4+
"encoding/json"
5+
"net/http"
6+
"os"
7+
"strconv"
8+
"time"
9+
10+
"github.com/google/uuid"
11+
"nhooyr.io/websocket"
12+
"nhooyr.io/websocket/wsjson"
13+
14+
"cdr.dev/slog"
15+
"github.com/coder/coder/coderd/database"
16+
"github.com/coder/coder/coderd/httpapi"
17+
"github.com/coder/coder/coderd/httpmw"
18+
"github.com/coder/coder/codersdk"
19+
)
20+
21+
const AgentStatIntervalEnv = "AGENT_STAT_INTERVAL"
22+
23+
func (api *API) workspaceAgentReportStats(rw http.ResponseWriter, r *http.Request) {
24+
api.websocketWaitMutex.Lock()
25+
api.websocketWaitGroup.Add(1)
26+
api.websocketWaitMutex.Unlock()
27+
defer api.websocketWaitGroup.Done()
28+
29+
workspaceAgent := httpmw.WorkspaceAgent(r)
30+
resource, err := api.Database.GetWorkspaceResourceByID(r.Context(), workspaceAgent.ResourceID)
31+
if err != nil {
32+
httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
33+
Message: "Failed to get workspace resource.",
34+
Detail: err.Error(),
35+
})
36+
return
37+
}
38+
39+
build, err := api.Database.GetWorkspaceBuildByJobID(r.Context(), resource.JobID)
40+
if err != nil {
41+
httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
42+
Message: "Failed to get build.",
43+
Detail: err.Error(),
44+
})
45+
return
46+
}
47+
48+
workspace, err := api.Database.GetWorkspaceByID(r.Context(), build.WorkspaceID)
49+
if err != nil {
50+
httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
51+
Message: "Failed to get workspace.",
52+
Detail: err.Error(),
53+
})
54+
return
55+
}
56+
57+
conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
58+
CompressionMode: websocket.CompressionDisabled,
59+
})
60+
if err != nil {
61+
httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
62+
Message: "Failed to accept websocket.",
63+
Detail: err.Error(),
64+
})
65+
return
66+
}
67+
defer conn.Close(websocket.StatusAbnormalClosure, "")
68+
69+
var interval = time.Minute
70+
71+
// Allow overriding the stat interval for debugging and testing purposes.
72+
intervalEnv, ok := os.LookupEnv(AgentStatIntervalEnv)
73+
if ok {
74+
intervalMs, err := strconv.Atoi(intervalEnv)
75+
if err != nil {
76+
api.Logger.Error(r.Context(), "parse agent stat interval",
77+
slog.F("interval", intervalEnv),
78+
slog.Error(err),
79+
)
80+
} else {
81+
interval = time.Millisecond * time.Duration(intervalMs)
82+
}
83+
}
84+
85+
ctx := r.Context()
86+
timer := time.NewTicker(interval)
87+
for {
88+
err := wsjson.Write(ctx, conn, codersdk.AgentStatsReportRequest{})
89+
if err != nil {
90+
httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
91+
Message: "Failed to write report request.",
92+
Detail: err.Error(),
93+
})
94+
return
95+
}
96+
var rep codersdk.AgentStatsReportResponse
97+
98+
err = wsjson.Read(ctx, conn, &rep)
99+
if err != nil {
100+
httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
101+
Message: "Failed to read report response.",
102+
Detail: err.Error(),
103+
})
104+
return
105+
}
106+
107+
repJSON, err := json.Marshal(rep)
108+
if err != nil {
109+
httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
110+
Message: "Failed to marshal stat json.",
111+
Detail: err.Error(),
112+
})
113+
return
114+
}
115+
116+
api.Logger.Debug(ctx, "read stats report",
117+
slog.F("agent", workspaceAgent.ID),
118+
slog.F("resource", resource.ID),
119+
slog.F("workspace", workspace.ID),
120+
slog.F("conns", rep.ProtocolStats),
121+
)
122+
_, err = api.Database.InsertAgentStat(ctx, database.InsertAgentStatParams{
123+
ID: uuid.NewString(),
124+
CreatedAt: time.Now(),
125+
AgentID: workspaceAgent.ID,
126+
WorkspaceID: build.WorkspaceID,
127+
UserID: workspace.OwnerID,
128+
Payload: json.RawMessage(repJSON),
129+
})
130+
if err != nil {
131+
httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
132+
Message: "Failed to insert agent stat.",
133+
Detail: err.Error(),
134+
})
135+
return
136+
}
137+
138+
select {
139+
case <-timer.C:
140+
continue
141+
case <-ctx.Done():
142+
conn.Close(websocket.StatusNormalClosure, "")
143+
return
144+
}
145+
}
146+
}

‎coderd/metrics_test.go

+94
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
package coderd_test
2+
3+
import (
4+
"context"
5+
"os"
6+
"testing"
7+
"time"
8+
9+
"github.com/google/uuid"
10+
"github.com/stretchr/testify/require"
11+
12+
"cdr.dev/slog/sloggers/slogtest"
13+
"github.com/coder/coder/agent"
14+
"github.com/coder/coder/coderd"
15+
"github.com/coder/coder/coderd/coderdtest"
16+
"github.com/coder/coder/codersdk"
17+
"github.com/coder/coder/peer"
18+
"github.com/coder/coder/provisioner/echo"
19+
"github.com/coder/coder/provisionersdk/proto"
20+
"github.com/coder/coder/testutil"
21+
)
22+
23+
func init() {
24+
os.Setenv(coderd.AgentStatIntervalEnv, "100")
25+
}
26+
func TestWorkspaceReportStats(t *testing.T) {
27+
t.Parallel()
28+
client := coderdtest.New(t, &coderdtest.Options{
29+
IncludeProvisionerD: true,
30+
})
31+
32+
user := coderdtest.CreateFirstUser(t, client)
33+
authToken := uuid.NewString()
34+
version := coderdtest.CreateTemplateVersion(t, client, user.OrganizationID, &echo.Responses{
35+
Parse: echo.ParseComplete,
36+
ProvisionDryRun: echo.ProvisionComplete,
37+
Provision: []*proto.Provision_Response{{
38+
Type: &proto.Provision_Response_Complete{
39+
Complete: &proto.Provision_Complete{
40+
Resources: []*proto.Resource{{
41+
Name: "example",
42+
Type: "aws_instance",
43+
Agents: []*proto.Agent{{
44+
Id: uuid.NewString(),
45+
Auth: &proto.Agent_Token{
46+
Token: authToken,
47+
},
48+
}},
49+
}},
50+
},
51+
},
52+
}},
53+
})
54+
template := coderdtest.CreateTemplate(t, client, user.OrganizationID, version.ID)
55+
coderdtest.AwaitTemplateVersionJob(t, client, version.ID)
56+
workspace := coderdtest.CreateWorkspace(t, client, user.OrganizationID, template.ID)
57+
coderdtest.AwaitWorkspaceBuildJob(t, client, workspace.LatestBuild.ID)
58+
59+
agentClient := codersdk.New(client.URL)
60+
agentClient.SessionToken = authToken
61+
agentCloser := agent.New(agentClient.ListenWorkspaceAgent, &agent.Options{
62+
Logger: slogtest.Make(t, nil),
63+
StatsReporter: agentClient.AgentReportStats,
64+
})
65+
defer func() {
66+
_ = agentCloser.Close()
67+
}()
68+
resources := coderdtest.AwaitWorkspaceAgents(t, client, workspace.LatestBuild.ID)
69+
70+
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
71+
defer cancel()
72+
73+
opts := &peer.ConnOptions{
74+
Logger: slogtest.Make(t, nil).Named("client"),
75+
}
76+
77+
conn, err := client.DialWorkspaceAgent(ctx, resources[0].Agents[0].ID, opts)
78+
require.NoError(t, err)
79+
defer func() {
80+
_ = conn.Close()
81+
}()
82+
83+
sshConn, err := conn.SSHClient()
84+
require.NoError(t, err)
85+
86+
session, err := sshConn.NewSession()
87+
require.NoError(t, err)
88+
89+
_, err = session.Output("echo hello")
90+
require.NoError(t, err)
91+
92+
time.Sleep(time.Second * 1)
93+
require.NoError(t, err)
94+
}

0 commit comments

Comments
 (0)