Skip to content

Commit 63fd494

Browse files
f0sself0ssel
authored
chore: watch workspace endpoint (coder#4060)
1 parent b340634 commit 63fd494

14 files changed

Lines changed: 682 additions & 339 deletions

File tree

‎coderd/database/databasefake/databasefake.go‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1431,6 +1431,22 @@ func (q *fakeQuerier) GetWorkspaceResourcesByJobID(_ context.Context, jobID uuid
14311431
return resources, nil
14321432
}
14331433

1434+
func (q *fakeQuerier) GetWorkspaceResourcesByJobIDs(_ context.Context, jobIDs []uuid.UUID) ([]database.WorkspaceResource, error) {
1435+
q.mutex.RLock()
1436+
defer q.mutex.RUnlock()
1437+
1438+
resources := make([]database.WorkspaceResource, 0)
1439+
for _, resource := range q.provisionerJobResources {
1440+
for _, jobID := range jobIDs {
1441+
if resource.JobID != jobID {
1442+
continue
1443+
}
1444+
resources = append(resources, resource)
1445+
}
1446+
}
1447+
return resources, nil
1448+
}
1449+
14341450
func (q *fakeQuerier) GetWorkspaceResourcesCreatedAfter(_ context.Context, after time.Time) ([]database.WorkspaceResource, error) {
14351451
q.mutex.RLock()
14361452
defer q.mutex.RUnlock()

‎coderd/database/querier.go‎

Lines changed: 1 addition & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries.sql.go‎

Lines changed: 41 additions & 0 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎coderd/database/queries/workspaceresources.sql‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,14 @@ FROM
1414
WHERE
1515
job_id = $1;
1616

17+
-- name: GetWorkspaceResourcesByJobIDs :many
18+
SELECT
19+
*
20+
FROM
21+
workspace_resources
22+
WHERE
23+
job_id = ANY(@ids :: uuid [ ]);
24+
1725
-- name: GetWorkspaceResourcesCreatedAfter :many
1826
SELECT * FROM workspace_resources WHERE created_at > $1;
1927

‎coderd/httpapi/httpapi.go‎

Lines changed: 76 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2,12 +2,16 @@ package httpapi
22

33
import (
44
"bytes"
5+
"context"
56
"encoding/json"
67
"errors"
78
"fmt"
9+
"io"
810
"net/http"
911
"reflect"
1012
"strings"
13+
"sync"
14+
"time"
1115

1216
"github.com/go-playground/validator/v10"
1317

@@ -144,3 +148,75 @@ func WebsocketCloseSprintf(format string, vars ...any) string {
144148

145149
return msg
146150
}
151+
152+
func ServerSentEventSender(rw http.ResponseWriter, r *http.Request) (func(ctx context.Context, sse codersdk.ServerSentEvent) error, error) {
153+
var mu sync.Mutex
154+
h := rw.Header()
155+
h.Set("Content-Type", "text/event-stream")
156+
h.Set("Cache-Control", "no-cache")
157+
h.Set("Connection", "keep-alive")
158+
h.Set("X-Accel-Buffering", "no")
159+
160+
f, ok := rw.(http.Flusher)
161+
if !ok {
162+
panic("http.ResponseWriter is not http.Flusher")
163+
}
164+
165+
// Send a heartbeat every 15 seconds to avoid the connection being killed.
166+
go func() {
167+
ticker := time.NewTicker(time.Second * 15)
168+
defer ticker.Stop()
169+
170+
for {
171+
select {
172+
case <-r.Context().Done():
173+
return
174+
case <-ticker.C:
175+
mu.Lock()
176+
_, err := io.WriteString(rw, fmt.Sprintf("event: %s\n\n", codersdk.ServerSentEventTypePing))
177+
if err != nil {
178+
mu.Unlock()
179+
return
180+
}
181+
f.Flush()
182+
mu.Unlock()
183+
}
184+
}
185+
}()
186+
187+
sendEvent := func(ctx context.Context, sse codersdk.ServerSentEvent) error {
188+
if ctx.Err() != nil {
189+
return ctx.Err()
190+
}
191+
192+
buf := &bytes.Buffer{}
193+
enc := json.NewEncoder(buf)
194+
195+
_, err := buf.Write([]byte(fmt.Sprintf("event: %s\ndata: ", sse.Type)))
196+
if err != nil {
197+
return err
198+
}
199+
200+
err = enc.Encode(sse.Data)
201+
if err != nil {
202+
return err
203+
}
204+
205+
err = buf.WriteByte('\n')
206+
if err != nil {
207+
return err
208+
}
209+
210+
mu.Lock()
211+
defer mu.Unlock()
212+
_, err = rw.Write(buf.Bytes())
213+
if err != nil {
214+
return err
215+
}
216+
f.Flush()
217+
218+
return nil
219+
}
220+
221+
return sendEvent, nil
222+
}

‎coderd/tracing/status_writer.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,3 +72,11 @@ func (w *StatusWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
7272
func (w *StatusWriter) ResponseBody() []byte {
7373
return w.responseBody
7474
}
75+
76+
func (w *StatusWriter) Flush() {
77+
f, ok := w.ResponseWriter.(http.Flusher)
78+
if !ok {
79+
panic("http.ResponseWriter is not http.Flusher")
80+
}
81+
f.Flush()
82+
}

0 commit comments

Comments
 (0)