move to channel over WorkspaceWatcher

f0ssel · f0ssel · commit 6e666be0d03b · 2022-05-18T16:25:07.000Z
diff --git a/coderd/workspaces_test.go b/coderd/workspaces_test.go
@@ -638,15 +638,14 @@ func TestWorkspaceWatcher(t *testing.T) {
 	w, err := client.Workspace(context.Background(), workspace.ID)
 	require.NoError(t, err)
 
-	ww, err := client.WatchWorkspace(context.Background(), w.ID)
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	wc, err := client.WatchWorkspace(ctx, w.ID)
 	require.NoError(t, err)
-	defer ww.Close()
-	for i := 0; i < 5; i++ {
-		_, err := ww.Read(context.Background())
-		require.NoError(t, err)
+	for i := 0; i < 3; i++ {
+		_, more := <-wc
+		require.True(t, more)
 	}
-	err = ww.Close()
-	require.NoError(t, err)
-	_, err = ww.Read(context.Background())
-	require.Error(t, err)
+	cancel()
+	require.EqualValues(t, codersdk.Workspace{}, <-wc)
 }
diff --git a/codersdk/workspaces.go b/codersdk/workspaces.go
@@ -94,38 +94,34 @@ func (c *Client) WorkspaceBuildByName(ctx context.Context, workspace uuid.UUID,
 	return workspaceBuild, json.NewDecoder(res.Body).Decode(&workspaceBuild)
 }
 
-type WorkspaceWatcher struct {
-	conn *websocket.Conn
-}
-
-func (w *WorkspaceWatcher) Read(ctx context.Context) (Workspace, error) {
-	var ws Workspace
-	err := wsjson.Read(ctx, w.conn, &ws)
-	if err != nil {
-		return ws, xerrors.Errorf("read workspace: %w", err)
-	}
-
-	return ws, nil
-}
-
-func (w *WorkspaceWatcher) Close() error {
-	err := w.conn.Close(websocket.StatusNormalClosure, "")
-	if err != nil {
-		return xerrors.Errorf("closing workspace watcher: %w", err)
-	}
-
-	return nil
-}
-
-func (c *Client) WatchWorkspace(ctx context.Context, id uuid.UUID) (*WorkspaceWatcher, error) {
+func (c *Client) WatchWorkspace(ctx context.Context, id uuid.UUID) (<-chan Workspace, error) {
 	conn, err := c.dialWebsocket(ctx, fmt.Sprintf("/api/v2/workspaces/%s/watch", id))
 	if err != nil {
 		return nil, err
 	}
+	wc := make(chan Workspace, 256)
+
+	go func() {
+		defer close(wc)
+		defer conn.Close(websocket.StatusNormalClosure, "")
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			default:
+				var ws Workspace
+				err := wsjson.Read(ctx, conn, &ws)
+				if err != nil {
+					conn.Close(websocket.StatusInternalError, "failed to read workspace")
+					return
+				}
+				wc <- ws
+			}
+		}
+	}()
 
-	return &WorkspaceWatcher{
-		conn: conn,
-	}, nil
+	return wc, nil
 }
 
 // UpdateWorkspaceAutostartRequest is a request to update a workspace's autostart schedule.
