coder · f0ssel · Sep 16, 2022 · Aug 25, 2022 · Aug 25, 2022 · Aug 25, 2022
diff --git a/coderd/database/databasefake/databasefake.go b/coderd/database/databasefake/databasefake.go
@@ -1431,6 +1431,22 @@ func (q *fakeQuerier) GetWorkspaceResourcesByJobID(_ context.Context, jobID uuid
 	return resources, nil
 }
 
+func (q *fakeQuerier) GetWorkspaceResourcesByJobIDs(_ context.Context, jobIDs []uuid.UUID) ([]database.WorkspaceResource, error) {
+	q.mutex.RLock()
+	defer q.mutex.RUnlock()
+
+	resources := make([]database.WorkspaceResource, 0)
+	for _, resource := range q.provisionerJobResources {
+		for _, jobID := range jobIDs {
+			if resource.JobID != jobID {
+				continue
+			}
+			resources = append(resources, resource)
+		}
+	}
+	return resources, nil
+}
+
 func (q *fakeQuerier) GetWorkspaceResourcesCreatedAfter(_ context.Context, after time.Time) ([]database.WorkspaceResource, error) {
 	q.mutex.RLock()
 	defer q.mutex.RUnlock()

diff --git a/coderd/database/querier.go b/coderd/database/querier.go
diff --git a/coderd/database/queries.sql.go b/coderd/database/queries.sql.go
diff --git a/coderd/database/queries/workspaceresources.sql b/coderd/database/queries/workspaceresources.sql
@@ -14,6 +14,14 @@ FROM
 WHERE
 	job_id = $1;
 
+-- name: GetWorkspaceResourcesByJobIDs :many
+SELECT
+	*
+FROM
+	workspace_resources
+WHERE
+	job_id = ANY(@ids :: uuid [ ]);
+
 -- name: GetWorkspaceResourcesCreatedAfter :many
 SELECT * FROM workspace_resources WHERE created_at > $1;
 

diff --git a/coderd/httpapi/httpapi.go b/coderd/httpapi/httpapi.go
@@ -2,12 +2,16 @@ package httpapi
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"errors"
 	"fmt"
+	"io"
 	"net/http"
 	"reflect"
 	"strings"
+	"sync"
+	"time"
 
 	"github.com/go-playground/validator/v10"
 
@@ -144,3 +148,75 @@ func WebsocketCloseSprintf(format string, vars ...any) string {
 
 	return msg
 }
+
+func ServerSentEventSender(rw http.ResponseWriter, r *http.Request) (func(ctx context.Context, sse codersdk.ServerSentEvent) error, error) {
+	var mu sync.Mutex
+	h := rw.Header()
+	h.Set("Content-Type", "text/event-stream")
+	h.Set("Cache-Control", "no-cache")
+	h.Set("Connection", "keep-alive")
+	h.Set("X-Accel-Buffering", "no")
+
+	f, ok := rw.(http.Flusher)
+	if !ok {
+		panic("http.ResponseWriter is not http.Flusher")
+	}
+
+	// Send a heartbeat every 15 seconds to avoid the connection being killed.
+	go func() {
+		ticker := time.NewTicker(time.Second * 15)
+		defer ticker.Stop()
+
+		for {
+			select {
+			case <-r.Context().Done():
+				return
+			case <-ticker.C:
+				mu.Lock()
+				_, err := io.WriteString(rw, fmt.Sprintf("event: %s\n\n", codersdk.ServerSentEventTypePing))
+				if err != nil {
+					mu.Unlock()
+					return
+				}
+				f.Flush()
+				mu.Unlock()
+			}
+		}
+	}()
+
+	sendEvent := func(ctx context.Context, sse codersdk.ServerSentEvent) error {
+		if ctx.Err() != nil {
+			return ctx.Err()
+		}
+
+		buf := &bytes.Buffer{}
+		enc := json.NewEncoder(buf)
+
+		_, err := buf.Write([]byte(fmt.Sprintf("event: %s\ndata: ", sse.Type)))
+		if err != nil {
+			return err
+		}
+
+		err = enc.Encode(sse.Data)
+		if err != nil {
+			return err
+		}
+
+		err = buf.WriteByte('\n')
+		if err != nil {
+			return err
+		}
+
+		mu.Lock()
+		defer mu.Unlock()
+		_, err = rw.Write(buf.Bytes())
+		if err != nil {
+			return err
+		}
+		f.Flush()
+
+		return nil
+	}
+
+	return sendEvent, nil
+}
diff --git a/coderd/tracing/status_writer.go b/coderd/tracing/status_writer.go
@@ -72,3 +72,11 @@ func (w *StatusWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
 func (w *StatusWriter) ResponseBody() []byte {
 	return w.responseBody
 }
+
+func (w *StatusWriter) Flush() {
+	f, ok := w.ResponseWriter.(http.Flusher)
+	if !ok {
+		panic("http.ResponseWriter is not http.Flusher")
+	}
+	f.Flush()
+}