Skip to content

Commit 99fe1bf

Browse files
mtojekmtojek
committed
Use errGroup
1 parent 4207dff commit 99fe1bf

File tree

4 files changed

+74
-32
lines changed

4 files changed

+74
-32
lines changed

‎cli/server.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -723,6 +723,10 @@ func (r *RootCmd) Server(newAPI func(context.Context, *coderd.Options) (*coderd.
723723
return xerrors.Errorf("register agent stats prometheus metric: %w", err)
724724
}
725725
defer closeAgentStatsFunc()
726+
727+
var metricsAggregator prometheusmetrics.MetricsAggregator
728+
options.UpdateAgentMetrics = metricsAggregator.Update
729+
options.PrometheusRegistry.MustRegister(&metricsAggregator)
726730
}
727731

728732
//nolint:revive

‎coderd/coderd.go

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,8 @@ import (
3838
"cdr.dev/slog"
3939

4040
"github.com/coder/coder/buildinfo"
41+
"github.com/coder/coder/codersdk/agentsdk"
42+
4143
// Used for swagger docs.
4244
_ "github.com/coder/coder/coderd/apidoc"
4345
"github.com/coder/coder/coderd/audit"
@@ -146,6 +148,8 @@ type Options struct {
146148
SSHConfig codersdk.SSHConfigResponse
147149

148150
HTTPClient *http.Client
151+
152+
UpdateAgentMetrics func(ctx context.Context, workspaceID uuid.UUID, agentID uuid.UUID, metrics []agentsdk.AgentMetric)
149153
}
150154

151155
// @title Coder API

‎coderd/prometheusmetrics/aggregator.go

Lines changed: 29 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,29 @@
1+
package prometheusmetrics
2+
3+
import (
4+
"context"
5+
"log"
6+
7+
"github.com/google/uuid"
8+
"github.com/prometheus/client_golang/prometheus"
9+
10+
"github.com/coder/coder/codersdk/agentsdk"
11+
)
12+
13+
type MetricsAggregator struct{}
14+
15+
var _ prometheus.Collector = new(MetricsAggregator)
16+
17+
// Describe function does not have any knowledge about the metrics schema,
18+
// so it does not emit anything.
19+
func (*MetricsAggregator) Describe(_ chan<- *prometheus.Desc) {
20+
}
21+
22+
func (ma *MetricsAggregator) Collect(ch chan<- prometheus.Metric) {
23+
}
24+
25+
// TODO Run function with done channel
26+
27+
func (ma *MetricsAggregator) Update(ctx context.Context, workspaceID uuid.UUID, agentID uuid.UUID, metrics []agentsdk.AgentMetric) {
28+
log.Printf("Workspace: %s, Agent: %s, Metrics: %v", workspaceID, agentID, metrics) // FIXME
29+
}

‎coderd/workspaceagents.go

Lines changed: 37 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,6 @@ import (
88
"errors"
99
"flag"
1010
"fmt"
11-
"log"
1211
"net"
1312
"net/http"
1413
"net/netip"
@@ -25,6 +24,7 @@ import (
2524
"github.com/google/uuid"
2625
"golang.org/x/exp/slices"
2726
"golang.org/x/mod/semver"
27+
"golang.org/x/sync/errgroup"
2828
"golang.org/x/xerrors"
2929
"nhooyr.io/websocket"
3030
"tailscale.com/tailcfg"
@@ -1214,42 +1214,47 @@ func (api *API) workspaceAgentReportStats(rw http.ResponseWriter, r *http.Reques
12141214
}
12151215

12161216
now := database.Now()
1217-
_, err = api.Database.InsertWorkspaceAgentStat(ctx, database.InsertWorkspaceAgentStatParams{
1218-
ID: uuid.New(),
1219-
CreatedAt: now,
1220-
AgentID: workspaceAgent.ID,
1221-
WorkspaceID: workspace.ID,
1222-
UserID: workspace.OwnerID,
1223-
TemplateID: workspace.TemplateID,
1224-
ConnectionsByProto: payload,
1225-
ConnectionCount: req.ConnectionCount,
1226-
RxPackets: req.RxPackets,
1227-
RxBytes: req.RxBytes,
1228-
TxPackets: req.TxPackets,
1229-
TxBytes: req.TxBytes,
1230-
SessionCountVSCode: req.SessionCountVSCode,
1231-
SessionCountJetBrains: req.SessionCountJetBrains,
1232-
SessionCountReconnectingPTY: req.SessionCountReconnectingPTY,
1233-
SessionCountSSH: req.SessionCountSSH,
1234-
ConnectionMedianLatencyMS: req.ConnectionMedianLatencyMS,
1235-
})
1236-
if err != nil {
1237-
httpapi.InternalServerError(rw, err)
1238-
return
1239-
}
12401217

1241-
if req.ConnectionCount > 0 {
1242-
err = api.Database.UpdateWorkspaceLastUsedAt(ctx, database.UpdateWorkspaceLastUsedAtParams{
1218+
var errGroup errgroup.Group
1219+
errGroup.Go(func() error {
1220+
_, err = api.Database.InsertWorkspaceAgentStat(ctx, database.InsertWorkspaceAgentStatParams{
1221+
ID: uuid.New(),
1222+
CreatedAt: now,
1223+
AgentID: workspaceAgent.ID,
1224+
WorkspaceID: workspace.ID,
1225+
UserID: workspace.OwnerID,
1226+
TemplateID: workspace.TemplateID,
1227+
ConnectionsByProto: payload,
1228+
ConnectionCount: req.ConnectionCount,
1229+
RxPackets: req.RxPackets,
1230+
RxBytes: req.RxBytes,
1231+
TxPackets: req.TxPackets,
1232+
TxBytes: req.TxBytes,
1233+
SessionCountVSCode: req.SessionCountVSCode,
1234+
SessionCountJetBrains: req.SessionCountJetBrains,
1235+
SessionCountReconnectingPTY: req.SessionCountReconnectingPTY,
1236+
SessionCountSSH: req.SessionCountSSH,
1237+
ConnectionMedianLatencyMS: req.ConnectionMedianLatencyMS,
1238+
})
1239+
return err
1240+
})
1241+
errGroup.Go(func() error {
1242+
return api.Database.UpdateWorkspaceLastUsedAt(ctx, database.UpdateWorkspaceLastUsedAtParams{
12431243
ID: workspace.ID,
12441244
LastUsedAt: now,
12451245
})
1246-
if err != nil {
1247-
httpapi.InternalServerError(rw, err)
1248-
return
1249-
}
1246+
})
1247+
if api.Options.UpdateAgentMetrics != nil {
1248+
errGroup.Go(func() error {
1249+
api.Options.UpdateAgentMetrics(ctx, workspace.ID, workspaceAgent.ID, req.Metrics)
1250+
return nil
1251+
})
1252+
}
1253+
err = errGroup.Wait()
1254+
if err != nil {
1255+
httpapi.InternalServerError(rw, err)
1256+
return
12501257
}
1251-
1252-
log.Println("Metrics: ", req.Metrics) // FIXME
12531258

12541259
httpapi.Write(ctx, rw, http.StatusOK, agentsdk.StatsResponse{
12551260
ReportInterval: api.AgentStatsRefreshInterval,

0 commit comments

Comments
 (0)