fix(codersdk): Wait for subscription in WatchWorkspaceAgentMetadata

mafredri · mafredri · commit 89f3db5e83a6 · 2023-06-13T12:14:12.000Z
diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go
@@ -1335,6 +1335,9 @@ func TestWorkspaceAgent_Metadata(t *testing.T) {
 	tooLongValueMetadata.CollectedAt = time.Now()
 	post("foo3", tooLongValueMetadata)
 	got := recvUpdate()[2]
+	for len(got.Result.Value) != maxValueLen {
+		got = recvUpdate()[2]
+	}
 	require.Len(t, got.Result.Value, maxValueLen)
 	require.NotEmpty(t, got.Result.Error)
 
diff --git a/codersdk/workspaceagents.go b/codersdk/workspaceagents.go
@@ -304,6 +304,7 @@ func (c *Client) WatchWorkspaceAgentMetadata(ctx context.Context, id uuid.UUID)
 
 	metadataChan := make(chan []WorkspaceAgentMetadata, 256)
 
+	ready := make(chan struct{})
 	watch := func() error {
 		res, err := c.Request(ctx, http.MethodGet, fmt.Sprintf("/api/v2/workspaceagents/%s/watch-metadata", id), nil)
 		if err != nil {
@@ -316,19 +317,24 @@ func (c *Client) WatchWorkspaceAgentMetadata(ctx context.Context, id uuid.UUID)
 		nextEvent := ServerSentEventReader(ctx, res.Body)
 		defer res.Body.Close()
 
+		firstEvent := true
 		for {
 			select {
 			case <-ctx.Done():
 				return ctx.Err()
 			default:
-				break
 			}
 
 			sse, err := nextEvent()
 			if err != nil {
 				return err
 			}
 
+			if firstEvent {
+				close(ready) // Only close ready after the first event is received.
+				firstEvent = false
+			}
+
 			b, ok := sse.Data.([]byte)
 			if !ok {
 				return xerrors.Errorf("unexpected data type: %T", sse.Data)
@@ -358,9 +364,18 @@ func (c *Client) WatchWorkspaceAgentMetadata(ctx context.Context, id uuid.UUID)
 	errorChan := make(chan error, 1)
 	go func() {
 		defer close(errorChan)
-		errorChan <- watch()
+		err := watch()
+		select {
+		case <-ready:
+		default:
+			close(ready) // Error before first event.
+		}
+		errorChan <- err
 	}()
 
+	// Wait until first event is received and the subscription is registered.
+	<-ready
+
 	return metadataChan, errorChan
 }
 
