chore: Add watch workspace endpoint

f0ssel · f0ssel · commit 4996df12a981 · 2022-05-17T20:38:05.000Z
diff --git a/coderd/coderd.go b/coderd/coderd.go
@@ -307,6 +307,7 @@ func New(options *Options) (http.Handler, func()) {
 			r.Route("/autostop", func(r chi.Router) {
 				r.Put("/", api.putWorkspaceAutostop)
 			})
+			r.HandleFunc("/watch", api.watchWorkspace)
 		})
 		r.Route("/workspacebuilds/{workspacebuild}", func(r chi.Router) {
 			r.Use(
diff --git a/coderd/workspaces.go b/coderd/workspaces.go
@@ -7,12 +7,16 @@ import (
 	"errors"
 	"fmt"
 	"net/http"
+	"time"
 
+	"cdr.dev/slog"
 	"github.com/go-chi/chi/v5"
 	"github.com/google/uuid"
 	"github.com/moby/moby/pkg/namesgenerator"
 	"golang.org/x/sync/errgroup"
 	"golang.org/x/xerrors"
+	"nhooyr.io/websocket"
+	"nhooyr.io/websocket/wsjson"
 
 	"github.com/coder/coder/coderd/autobuild/schedule"
 	"github.com/coder/coder/coderd/database"
@@ -496,6 +500,94 @@ func (api *api) putWorkspaceAutostop(rw http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func (api *api) watchWorkspace(rw http.ResponseWriter, r *http.Request) {
+	workspace := httpmw.WorkspaceParam(r)
+
+	c, err := websocket.Accept(rw, r, &websocket.AcceptOptions{
+		// Fix for Safari 15.1:
+		// There is a bug in latest Safari in which compressed web socket traffic
+		// isn't handled correctly. Turning off compression is a workaround:
+		// https://github.com/nhooyr/websocket/issues/218
+		CompressionMode: websocket.CompressionDisabled,
+	})
+	if err != nil {
+		api.Logger.Warn(r.Context(), "accept websocket connection", slog.Error(err))
+		return
+	}
+	defer c.Close(websocket.StatusInternalError, "internal error")
+
+	ctx := c.CloseRead(r.Context())
+
+	// Send a heartbeat every 15 seconds to avoid the websocket being killed.
+	go func() {
+		ticker := time.NewTicker(time.Second * 15)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-ticker.C:
+				err := c.Ping(ctx)
+				if err != nil {
+					return
+				}
+			}
+		}
+	}()
+
+	t := time.NewTicker(time.Second * 1)
+	defer t.Stop()
+	for {
+		select {
+		case <-t.C:
+			workspace, err := api.Database.GetWorkspaceByID(r.Context(), workspace.ID)
+			if err != nil {
+				_ = wsjson.Write(ctx, c, httpapi.Response{
+					Message: fmt.Sprintf("get workspace: %s", err),
+				})
+				return
+			}
+			build, err := api.Database.GetWorkspaceBuildByWorkspaceIDWithoutAfter(r.Context(), workspace.ID)
+			if err != nil {
+				_ = wsjson.Write(ctx, c, httpapi.Response{
+					Message: fmt.Sprintf("get workspace build: %s", err),
+				})
+				return
+			}
+			var (
+				group    errgroup.Group
+				job      database.ProvisionerJob
+				template database.Template
+				owner    database.User
+			)
+			group.Go(func() (err error) {
+				job, err = api.Database.GetProvisionerJobByID(r.Context(), build.JobID)
+				return err
+			})
+			group.Go(func() (err error) {
+				template, err = api.Database.GetTemplateByID(r.Context(), workspace.TemplateID)
+				return err
+			})
+			group.Go(func() (err error) {
+				owner, err = api.Database.GetUserByID(r.Context(), workspace.OwnerID)
+				return err
+			})
+			err = group.Wait()
+			if err != nil {
+				_ = wsjson.Write(ctx, c, httpapi.Response{
+					Message: fmt.Sprintf("fetch resource: %s", err),
+				})
+				return
+			}
+
+			_ = wsjson.Write(ctx, c, convertWorkspace(workspace, convertWorkspaceBuild(build, convertProvisionerJob(job)), template, owner))
+		case <-ctx.Done():
+			return
+		}
+	}
+}
+
 func convertWorkspaces(ctx context.Context, db database.Store, workspaces []database.Workspace) ([]codersdk.Workspace, error) {
 	workspaceIDs := make([]uuid.UUID, 0, len(workspaces))
 	templateIDs := make([]uuid.UUID, 0, len(workspaces))
