fix: handle new agent stat format correctly (#14576)

---------

Co-authored-by: Ethan Dickson <ethan@coder.com>
This commit is contained in:
Garrett Delfosse
2024-09-20 01:52:14 +10:00
committed by GitHub
co-authored by Ethan Dickson
parent 37885e2e82
commit 922f4c545f
28 changed files with 1495 additions and 70 deletions
+5 -1
View File
@@ -25,7 +25,7 @@ const (
)
type Batcher interface {
Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error
Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats, usage bool) error
}
// DBBatcher holds a buffer of agent stats and periodically flushes them to
@@ -138,6 +138,7 @@ func (b *DBBatcher) Add(
userID uuid.UUID,
workspaceID uuid.UUID,
st *agentproto.Stats,
usage bool,
) error {
b.mu.Lock()
defer b.mu.Unlock()
@@ -165,6 +166,7 @@ func (b *DBBatcher) Add(
b.buf.SessionCountReconnectingPTY = append(b.buf.SessionCountReconnectingPTY, st.SessionCountReconnectingPty)
b.buf.SessionCountSSH = append(b.buf.SessionCountSSH, st.SessionCountSsh)
b.buf.ConnectionMedianLatencyMS = append(b.buf.ConnectionMedianLatencyMS, st.ConnectionMedianLatencyMs)
b.buf.Usage = append(b.buf.Usage, usage)
// If the buffer is over 80% full, signal the flusher to flush immediately.
// We want to trigger flushes early to reduce the likelihood of
@@ -279,6 +281,7 @@ func (b *DBBatcher) initBuf(size int) {
SessionCountReconnectingPTY: make([]int64, 0, b.batchSize),
SessionCountSSH: make([]int64, 0, b.batchSize),
ConnectionMedianLatencyMS: make([]float64, 0, b.batchSize),
Usage: make([]bool, 0, b.batchSize),
}
b.connectionsByProto = make([]map[string]int64, 0, size)
@@ -302,5 +305,6 @@ func (b *DBBatcher) resetBuf() {
b.buf.SessionCountReconnectingPTY = b.buf.SessionCountReconnectingPTY[:0]
b.buf.SessionCountSSH = b.buf.SessionCountSSH[:0]
b.buf.ConnectionMedianLatencyMS = b.buf.ConnectionMedianLatencyMS[:0]
b.buf.Usage = b.buf.Usage[:0]
b.connectionsByProto = b.connectionsByProto[:0]
}
@@ -63,7 +63,7 @@ func TestBatchStats(t *testing.T) {
// Given: a single data point is added for workspace
t2 := t1.Add(time.Second)
t.Logf("inserting 1 stat")
require.NoError(t, b.Add(t2.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t)))
require.NoError(t, b.Add(t2.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false))
// When: it becomes time to report stats
// Signal a tick and wait for a flush to complete.
@@ -87,9 +87,9 @@ func TestBatchStats(t *testing.T) {
t.Logf("inserting %d stats", defaultBufferSize)
for i := 0; i < defaultBufferSize; i++ {
if i%2 == 0 {
require.NoError(t, b.Add(t3.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t)))
require.NoError(t, b.Add(t3.Add(time.Millisecond), deps1.Agent.ID, deps1.User.ID, deps1.Template.ID, deps1.Workspace.ID, randStats(t), false))
} else {
require.NoError(t, b.Add(t3.Add(time.Millisecond), deps2.Agent.ID, deps2.User.ID, deps2.Template.ID, deps2.Workspace.ID, randStats(t)))
require.NoError(t, b.Add(t3.Add(time.Millisecond), deps2.Agent.ID, deps2.User.ID, deps2.Template.ID, deps2.Workspace.ID, randStats(t), false))
}
}
}()
+2 -2
View File
@@ -118,7 +118,7 @@ func (r *Reporter) ReportAppStats(ctx context.Context, stats []workspaceapps.Sta
return nil
}
func (r *Reporter) ReportAgentStats(ctx context.Context, now time.Time, workspace database.Workspace, workspaceAgent database.WorkspaceAgent, templateName string, stats *agentproto.Stats) error {
func (r *Reporter) ReportAgentStats(ctx context.Context, now time.Time, workspace database.Workspace, workspaceAgent database.WorkspaceAgent, templateName string, stats *agentproto.Stats, usage bool) error {
if stats.ConnectionCount > 0 {
var nextAutostart time.Time
if workspace.AutostartSchedule.String != "" {
@@ -143,7 +143,7 @@ func (r *Reporter) ReportAgentStats(ctx context.Context, now time.Time, workspac
var errGroup errgroup.Group
errGroup.Go(func() error {
err := r.opts.StatsBatcher.Add(now, workspaceAgent.ID, workspace.TemplateID, workspace.OwnerID, workspace.ID, stats)
err := r.opts.StatsBatcher.Add(now, workspaceAgent.ID, workspace.TemplateID, workspace.OwnerID, workspace.ID, stats, usage)
if err != nil {
r.opts.Logger.Error(ctx, "add agent stats to batcher", slog.Error(err))
return xerrors.Errorf("insert workspace agent stats batch: %w", err)
@@ -20,11 +20,12 @@ type StatsBatcher struct {
LastUserID uuid.UUID
LastWorkspaceID uuid.UUID
LastStats *agentproto.Stats
LastUsage bool
}
var _ workspacestats.Batcher = &StatsBatcher{}
func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats) error {
func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUID, userID uuid.UUID, workspaceID uuid.UUID, st *agentproto.Stats, usage bool) error {
b.Mu.Lock()
defer b.Mu.Unlock()
b.Called++
@@ -34,5 +35,6 @@ func (b *StatsBatcher) Add(now time.Time, agentID uuid.UUID, templateID uuid.UUI
b.LastUserID = userID
b.LastWorkspaceID = workspaceID
b.LastStats = st
b.LastUsage = usage
return nil
}