Add workspace report stats endpoint

ammario · ammario · commit bdde5505caaf · 2022-08-29T20:55:49.000Z
diff --git a/agent/stats.go b/agent/stats.go
@@ -16,8 +16,7 @@ type ConnStats struct {
 	RxBytes uint64 `json:"rx_bytes,omitempty"`
 
 	// TxBytes must be read with atomic.
-	TxBytes uint64 `json:"tx_bytes,omitempty"`
-
+	TxBytes  uint64 `json:"tx_bytes,omitempty"`
 	net.Conn `json:"-"`
 }
 
diff --git a/coderd/workspaceagents.go b/coderd/workspaceagents.go
@@ -8,7 +8,9 @@ import (
 	"io"
 	"net"
 	"net/http"
+	"os"
 	"strconv"
+	"strings"
 	"time"
 
 	"github.com/google/uuid"
@@ -17,6 +19,7 @@ import (
 	"golang.org/x/xerrors"
 	"inet.af/netaddr"
 	"nhooyr.io/websocket"
+	"nhooyr.io/websocket/wsjson"
 	"tailscale.com/types/key"
 
 	"cdr.dev/slog"
@@ -152,6 +155,80 @@ func (api *API) workspaceAgentMetadata(rw http.ResponseWriter, r *http.Request)
 	})
 }
 
+var inTest = strings.HasSuffix(os.Args[0], ".test")
+
+func (api *API) workspaceAgentReportStats(rw http.ResponseWriter, r *http.Request) {
+	api.websocketWaitMutex.Lock()
+	api.websocketWaitGroup.Add(1)
+	api.websocketWaitMutex.Unlock()
+	defer api.websocketWaitGroup.Done()
+
+	workspaceAgent := httpmw.WorkspaceAgent(r)
+	workspace, err := api.Database.GetWorkspaceResourceByID(r.Context(), workspaceAgent.ResourceID)
+	if err != nil {
+		httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Failed to accept websocket.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+
+	conn, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
+		CompressionMode: websocket.CompressionDisabled,
+	})
+	if err != nil {
+		httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
+			Message: "Failed to accept websocket.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	defer conn.Close(websocket.StatusAbnormalClosure, "")
+
+	var interval = time.Minute
+	if inTest {
+		interval = 100 * time.Millisecond
+	}
+
+	ctx := r.Context()
+	timer := time.NewTimer(interval)
+	for {
+		err := wsjson.Write(ctx, conn, codersdk.AgentStatsReportRequest{})
+		if err != nil {
+			httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
+				Message: "Failed to write report request.",
+				Detail:  err.Error(),
+			})
+			return
+		}
+		var rep codersdk.AgentStatsReportResponse
+
+		err = wsjson.Read(ctx, conn, &rep)
+		if err != nil {
+			httpapi.Write(rw, http.StatusBadRequest, codersdk.Response{
+				Message: "Failed to read report response.",
+				Detail:  err.Error(),
+			})
+			return
+		}
+
+		api.Logger.Debug(ctx, "read stats report",
+			slog.F("agent", workspaceAgent.ID),
+			slog.F("workspace", workspace.ID),
+			slog.F("report", rep),
+		)
+
+		select {
+		case <-timer.C:
+			continue
+		case <-ctx.Done():
+			conn.Close(websocket.StatusNormalClosure, "")
+			return
+		}
+	}
+
+}
+
 func (api *API) workspaceAgentListen(rw http.ResponseWriter, r *http.Request) {
 	api.websocketWaitMutex.Lock()
 	api.websocketWaitGroup.Add(1)
diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go
@@ -480,3 +480,14 @@ func (c *Client) turnProxyDialer(ctx context.Context, httpClient *http.Client, p
 		return websocket.NetConn(ctx, conn, websocket.MessageBinary), nil
 	})
 }
+
+// AgentStatsReportRequest is a WebSocket request by coderd
+// to the agent for stats.
+type AgentStatsReportRequest struct {
+}
+
+// AgentStatsReportResponse is returned for each report
+// request by the agent.
+type AgentStatsReportResponse struct {
+	Conns []agent.ConnStats `json:"conns,omitempty"`
+}

Original file line number	Diff line number	Diff line change
`@@ -16,8 +16,7 @@ type ConnStats struct {`
`16`	`16`	RxBytes uint64 `json:"rx_bytes,omitempty"`
`17`	`17`
`18`	`18`	`// TxBytes must be read with atomic.`
`19`		- TxBytes uint64 `json:"tx_bytes,omitempty"`
`20`		`-`
	`19`	+ TxBytes uint64 `json:"tx_bytes,omitempty"`
`21`	`20`	net.Conn `json:"-"`
`22`	`21`	`}`
`23`	`22`