fix(coderd): fix memory leak in watchWorkspaceAgentMetadata (#10685)

Fixes #10550

coderd/workspaceagents.go

@@ -1785,10 +1785,10 @@ func (api *API) workspaceAgentUpdateMetadata(ctx context.Context, workspaceAgent
 
 	datum := database.UpdateWorkspaceAgentMetadataParams{
 		WorkspaceAgentID: workspaceAgent.ID,
-		Key:              []string{},
-		Value:            []string{},
-		Error:            []string{},
-		CollectedAt:      []time.Time{},
+		Key:              make([]string, 0, len(req.Metadata)),
+		Value:            make([]string, 0, len(req.Metadata)),
+		Error:            make([]string, 0, len(req.Metadata)),
+		CollectedAt:      make([]time.Time, 0, len(req.Metadata)),
 	}
 
 	for _, md := range req.Metadata {
@@ -1853,22 +1853,28 @@ func (api *API) workspaceAgentUpdateMetadata(ctx context.Context, workspaceAgent
 // @Router /workspaceagents/{workspaceagent}/watch-metadata [get]
 // @x-apidocgen {"skip": true}
 func (api *API) watchWorkspaceAgentMetadata(rw http.ResponseWriter, r *http.Request) {
-	var (
-		ctx            = r.Context()
-		workspaceAgent = httpmw.WorkspaceAgentParam(r)
-		log            = api.Logger.Named("workspace_metadata_watcher").With(
-			slog.F("workspace_agent_id", workspaceAgent.ID),
-		)
+	// Allow us to interrupt watch via cancel.
+	ctx, cancel := context.WithCancel(r.Context())
+	defer cancel()
+	r = r.WithContext(ctx) // Rewire context for SSE cancellation.
+
+	workspaceAgent := httpmw.WorkspaceAgentParam(r)
+	log := api.Logger.Named("workspace_metadata_watcher").With(
+		slog.F("workspace_agent_id", workspaceAgent.ID),
 	)
 
 	// Send metadata on updates, we must ensure subscription before sending
 	// initial metadata to guarantee that events in-between are not missed.
 	update := make(chan workspaceAgentMetadataChannelPayload, 1)
 	cancelSub, err := api.Pubsub.Subscribe(watchWorkspaceAgentMetadataChannel(workspaceAgent.ID), func(_ context.Context, byt []byte) {
+		if ctx.Err() != nil {
+			return
+		}
+
 		var payload workspaceAgentMetadataChannelPayload
 		err := json.Unmarshal(byt, &payload)
 		if err != nil {
-			api.Logger.Error(ctx, "failed to unmarshal pubsub message", slog.Error(err))
+			log.Error(ctx, "failed to unmarshal pubsub message", slog.Error(err))
 			return
 		}
 
@@ -1876,18 +1882,7 @@ func (api *API) watchWorkspaceAgentMetadata(rw http.ResponseWriter, r *http.Requ
 
 		select {
 		case prev := <-update:
-			// This update wasn't consumed yet, merge the keys.
-			newKeysSet := make(map[string]struct{})
-			for _, key := range payload.Keys {
-				newKeysSet[key] = struct{}{}
-			}
-			keys := prev.Keys
-			for _, key := range prev.Keys {
-				if _, ok := newKeysSet[key]; !ok {
-					keys = append(keys, key)
-				}
-			}
-			payload.Keys = keys
+			payload.Keys = appendUnique(prev.Keys, payload.Keys)
 		default:
 		}
 		// This can never block since we pop and merge beforehand.
@@ -1899,22 +1894,9 @@ func (api *API) watchWorkspaceAgentMetadata(rw http.ResponseWriter, r *http.Requ
 	}
 	defer cancelSub()
 
-	sseSendEvent, sseSenderClosed, err := httpapi.ServerSentEventSender(rw, r)
-	if err != nil {
-		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
-			Message: "Internal error setting up server-sent events.",
-			Detail:  err.Error(),
-		})
-		return
-	}
-	// Prevent handler from returning until the sender is closed.
-	defer func() {
-		<-sseSenderClosed
-	}()
-
 	// We always use the original Request context because it contains
 	// the RBAC actor.
-	md, err := api.Database.GetWorkspaceAgentMetadata(ctx, database.GetWorkspaceAgentMetadataParams{
+	initialMD, err := api.Database.GetWorkspaceAgentMetadata(ctx, database.GetWorkspaceAgentMetadataParams{
 		WorkspaceAgentID: workspaceAgent.ID,
 		Keys:             nil,
 	})
@@ -1926,15 +1908,45 @@ func (api *API) watchWorkspaceAgentMetadata(rw http.ResponseWriter, r *http.Requ
 		return
 	}
 
-	metadataMap := make(map[string]database.WorkspaceAgentMetadatum)
-	for _, datum := range md {
+	log.Debug(ctx, "got initial metadata", "num", len(initialMD))
+
+	metadataMap := make(map[string]database.WorkspaceAgentMetadatum, len(initialMD))
+	for _, datum := range initialMD {
 		metadataMap[datum.Key] = datum
 	}
+	//nolint:ineffassign // Release memory.
+	initialMD = nil
+
+	sseSendEvent, sseSenderClosed, err := httpapi.ServerSentEventSender(rw, r)
+	if err != nil {
+		httpapi.Write(ctx, rw, http.StatusInternalServerError, codersdk.Response{
+			Message: "Internal error setting up server-sent events.",
+			Detail:  err.Error(),
+		})
+		return
+	}
+	// Prevent handler from returning until the sender is closed.
+	defer func() {
+		cancel()
+		<-sseSenderClosed
+	}()
+	// Synchronize cancellation from SSE -> context, this lets us simplify the
+	// cancellation logic.
+	go func() {
+		select {
+		case <-ctx.Done():
+		case <-sseSenderClosed:
+			cancel()
+		}
+	}()
 
 	var lastSend time.Time
 	sendMetadata := func() {
 		lastSend = time.Now()
 		values := maps.Values(metadataMap)
+
+		log.Debug(ctx, "sending metadata", "num", len(values))
+
 		_ = sseSendEvent(ctx, codersdk.ServerSentEvent{
 			Type: codersdk.ServerSentEventTypeData,
 			Data: convertWorkspaceAgentMetadata(values),
@@ -1953,10 +1965,11 @@ func (api *API) watchWorkspaceAgentMetadata(rw http.ResponseWriter, r *http.Requ
 	fetchedMetadata := make(chan []database.WorkspaceAgentMetadatum)
 	go func() {
 		defer close(fetchedMetadata)
+		defer cancel()
 
 		for {
 			select {
-			case <-sseSenderClosed:
+			case <-ctx.Done():
 				return
 			case payload := <-update:
 				md, err := api.Database.GetWorkspaceAgentMetadata(ctx, database.GetWorkspaceAgentMetadataParams{
@@ -1966,24 +1979,35 @@ func (api *API) watchWorkspaceAgentMetadata(rw http.ResponseWriter, r *http.Requ
 				if err != nil {
 					if !errors.Is(err, context.Canceled) {
 						log.Error(ctx, "failed to get metadata", slog.Error(err))
+						_ = sseSendEvent(ctx, codersdk.ServerSentEvent{
+							Type: codersdk.ServerSentEventTypeError,
+							Data: codersdk.Response{
+								Message: "Failed to get metadata.",
+								Detail:  err.Error(),
+							},
+						})
 					}
 					return
 				}
 				select {
-				case <-sseSenderClosed:
+				case <-ctx.Done():
 					return
 				// We want to block here to avoid constantly pinging the
 				// database when the metadata isn't being processed.
 				case fetchedMetadata <- md:
+					log.Debug(ctx, "fetched metadata update for keys", "keys", payload.Keys, "num", len(md))
 				}
 			}
 		}
 	}()
+	defer func() {
+		<-fetchedMetadata
+	}()
 
 	pendingChanges := true
 	for {
 		select {
-		case <-sseSenderClosed:
+		case <-ctx.Done():
 			return
 		case md, ok := <-fetchedMetadata:
 			if !ok {
@@ -2007,9 +2031,24 @@ func (api *API) watchWorkspaceAgentMetadata(rw http.ResponseWriter, r *http.Requ
 	}
 }
 
+// appendUnique is like append and adds elements from src to dst,
+// skipping any elements that already exist in dst.
+func appendUnique[T comparable](dst, src []T) []T {
+	exists := make(map[T]struct{}, len(dst))
+	for _, key := range dst {
+		exists[key] = struct{}{}
+	}
+	for _, key := range src {
+		if _, ok := exists[key]; !ok {
+			dst = append(dst, key)
+		}
+	}
+	return dst
+}
+
 func convertWorkspaceAgentMetadata(db []database.WorkspaceAgentMetadatum) []codersdk.WorkspaceAgentMetadata {
 	// An empty array is easier for clients to handle than a null.
-	result := []codersdk.WorkspaceAgentMetadata{}
+	result := make([]codersdk.WorkspaceAgentMetadata, 0, len(db))
 	for _, datum := range db {
 		result = append(result, codersdk.WorkspaceAgentMetadata{
 			Result: codersdk.WorkspaceAgentMetadataResult{