diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index a0c82fa..13a9293 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -44,6 +44,8 @@ The worker Pods only have access to the queue, file service, and squid proxy. Th The control microservice is the server that receives requests from the user. It does create the corresponding workers, sends the messages to the queue, and responds to the user the response payload. Also it does store and serve the user snippets from Etcd. +A run is started with `POST /service/control/run`. Clients that send `Accept: text/event-stream` receive `202 { id }` immediately and then attach to `GET /service/control/run/:id/log-watch` (SSE) for live stdout/stderr. Other clients wait for the same JSON payload as before (`success`, `output`, `files`, …). Workers still publish over RabbitMQ: `log` events while the process runs, then a final `done` payload. Each worker pod still executes a single job and is replaced. + ### Worker For each of the languages, there are individual Docker images and worker implementations since each language gets executed differently. diff --git a/control-service/handle_run_correlation_test.go b/control-service/handle_run_correlation_test.go index bd555fa..d3eaf0f 100644 --- a/control-service/handle_run_correlation_test.go +++ b/control-service/handle_run_correlation_test.go @@ -57,6 +57,7 @@ func newRunTestServer(t *testing.T) *httptest.Server { workers: map[workertypes.WorkerLanguage]*Workers{ workertypes.WorkerLanguageJavaScript: {workers: make(chan *Worker)}, }, + runs: newRunHub(), } s.initializeHttpServer() return httptest.NewServer(s.echo) diff --git a/control-service/main.go b/control-service/main.go index 53cadc5..5885ee6 100644 --- a/control-service/main.go +++ b/control-service/main.go @@ -57,6 +57,7 @@ type server struct { amqpErrorChan chan *amqp.Error workers map[workertypes.WorkerLanguage]*Workers + runs *runHub } func newServer() (*server, error) { @@ -121,6 +122,7 @@ func newServer() (*server, error) { amqpConnection: amqpConnection, amqpErrorChan: amqpErrorChan, workers: workersMap, + runs: newRunHub(), } s.initializeHttpServer() @@ -134,6 +136,7 @@ func (s *server) initializeHttpServer() { s.echo.GET("/service/control/health", s.handleHealth) s.echo.HEAD("/service/control/health", s.handleHealth) s.echo.POST("/service/control/run", s.handleRun) + s.echo.GET("/service/control/run/:id/log-watch", s.handleLogWatch) s.echo.GET("/service/control/share/get/:id", s.handleShareGet) s.echo.POST("/service/control/share/create", s.handleShareCreate) } @@ -216,58 +219,74 @@ func (s *server) handleRun(c *echo.Context) error { logger.Infof("Received code: '%s'", req.Code) logger.Info("Obtained worker successfully") logger.Info("Publishing job") + session := newRunSession(requestID) + s.runs.Put(requestID, session) + workers.replies.Store(worker.id, session) if err := worker.Publish(req.Code, req.RequestID, req.TestID); err != nil { logger.Errorf("could not create new worker job: %v", err) + s.runs.Delete(requestID) + workers.replies.Delete(worker.id) return respondError(c, http.StatusInternalServerError, requestID, testID, logBuffer, "could not create new worker job") } logger.Println("Published message") - start := time.Now() - - var payload *workertypes.WorkerResponsePayload - timeout := false - select { - case payload = <-worker.Subscribe(): - payload.Duration = time.Since(start).Milliseconds() - logger.Println("Received response successfully") - case <-time.After(EXECUTION_TIMEOUT * time.Second): - logger.Println("Got execution timeout!") - timeout = true - } - go func() { + select { + case <-session.finished: + logger.Println("Received response successfully") + case <-time.After(EXECUTION_TIMEOUT * time.Second): + logger.Println("Got execution timeout!") + session.Fail("Execution timeout!") + } + logger.Println("Starting worker cleanup") if err := worker.Cleanup(); err != nil { logger.Printf("could not cleanup worker: %v", err) - return + } else { + logger.Println("Finished worker cleanup") } - logger.Println("Finished worker cleanup") logger.Println("Adding new worker") if err := workers.AddWorkers(1); err != nil { logger.Printf("could not create new worker: %v", err) - return + } else { + logger.Println("Added new worker successfully") } - logger.Println("Added new worker successfully") - }() - if timeout { - return c.JSON(http.StatusServiceUnavailable, map[string]any{ - "error": "Execution timeout!", - "requestId": requestID, - "testId": testID, - "logs": map[string]any{ - "control": logBuffer.String(), - }, + time.AfterFunc(runSessionTTL, func() { + s.runs.Delete(requestID) }) - } + }() - payload.RequestID = requestID - payload.TestID = testID - if !payload.Success { - return c.JSON(http.StatusBadRequest, payload) + if wantsJSONWait(c.Request().Header.Get("Accept")) { + <-session.finished + payload, timedOut := session.Result() + if timedOut || (payload != nil && payload.Error == "Execution timeout!") { + return c.JSON(http.StatusServiceUnavailable, map[string]any{ + "error": "Execution timeout!", + "requestId": requestID, + "testId": testID, + "logs": map[string]any{ + "control": logBuffer.String(), + }, + }) + } + if payload == nil { + return respondError(c, http.StatusInternalServerError, requestID, testID, logBuffer, "missing worker response") + } + payload.RequestID = requestID + payload.TestID = testID + if !payload.Success { + return c.JSON(http.StatusBadRequest, payload) + } + return c.JSON(http.StatusOK, payload) } - return c.JSON(http.StatusOK, payload) + + return c.JSON(http.StatusAccepted, map[string]any{ + "id": requestID, + "requestId": requestID, + "testId": testID, + }) } func (s *server) handleShareGet(c *echo.Context) error { diff --git a/control-service/runs.go b/control-service/runs.go new file mode 100644 index 0000000..0cff532 --- /dev/null +++ b/control-service/runs.go @@ -0,0 +1,283 @@ +package main + +import ( + "encoding/json" + "fmt" + "net/http" + "strings" + "sync" + "time" + + "github.com/labstack/echo/v5" + "github.com/mxschmitt/try-playwright/internal/workertypes" + log "github.com/sirupsen/logrus" +) + +const runSessionTTL = 2 * time.Minute + +type runEvent struct { + Type string + Line string + Done *workertypes.WorkerResponsePayload + Error string +} + +type runSession struct { + id string + start time.Time + mu sync.Mutex + logs []string + done *workertypes.WorkerResponsePayload + timeout bool + fail string + subs []chan runEvent + finished chan struct{} + once sync.Once +} + +func newRunSession(id string) *runSession { + return &runSession{ + id: id, + start: time.Now(), + finished: make(chan struct{}), + } +} + +func (s *runSession) markFinished() { + s.once.Do(func() { close(s.finished) }) +} + +func (s *runSession) AppendLog(line string) { + ev := runEvent{Type: workertypes.WorkerEventLog, Line: line} + s.mu.Lock() + s.logs = append(s.logs, line) + subs := append([]chan runEvent(nil), s.subs...) + s.mu.Unlock() + for _, ch := range subs { + select { + case ch <- ev: + default: + } + } +} + +func (s *runSession) Complete(payload *workertypes.WorkerResponsePayload) { + if payload == nil { + payload = &workertypes.WorkerResponsePayload{} + } + payload.Duration = time.Since(s.start).Milliseconds() + payload.RequestID = s.id + ev := runEvent{Type: workertypes.WorkerEventDone, Done: payload} + s.mu.Lock() + if s.done != nil || s.fail != "" { + s.mu.Unlock() + return + } + s.done = payload + subs := s.subs + s.subs = nil + s.mu.Unlock() + for _, ch := range subs { + ch <- ev + close(ch) + } + s.markFinished() +} + +func (s *runSession) Fail(msg string) { + payload := &workertypes.WorkerResponsePayload{ + Success: false, + Error: msg, + RequestID: s.id, + Files: []workertypes.File{}, + } + payload.Duration = time.Since(s.start).Milliseconds() + ev := runEvent{Type: "error", Error: msg, Done: payload} + s.mu.Lock() + if s.done != nil || s.fail != "" { + s.mu.Unlock() + return + } + s.fail = msg + s.timeout = msg == "Execution timeout!" + s.done = payload + subs := s.subs + s.subs = nil + s.mu.Unlock() + for _, ch := range subs { + ch <- ev + close(ch) + } + s.markFinished() +} + +func (s *runSession) Subscribe() (<-chan runEvent, func()) { + ch := make(chan runEvent, 256) + s.mu.Lock() + defer s.mu.Unlock() + for _, line := range s.logs { + ch <- runEvent{Type: workertypes.WorkerEventLog, Line: line} + } + if s.done != nil { + if s.fail != "" { + ch <- runEvent{Type: "error", Error: s.fail, Done: s.done} + } else { + ch <- runEvent{Type: workertypes.WorkerEventDone, Done: s.done} + } + close(ch) + return ch, func() {} + } + s.subs = append(s.subs, ch) + return ch, func() { + s.mu.Lock() + defer s.mu.Unlock() + filtered := s.subs[:0] + for _, existing := range s.subs { + if existing != ch { + filtered = append(filtered, existing) + } + } + s.subs = filtered + select { + case <-ch: + default: + } + } +} + +func (s *runSession) Result() (*workertypes.WorkerResponsePayload, bool) { + s.mu.Lock() + defer s.mu.Unlock() + return s.done, s.timeout +} + +type runHub struct { + mu sync.Mutex + sessions map[string]*runSession +} + +func newRunHub() *runHub { + return &runHub{sessions: map[string]*runSession{}} +} + +func (h *runHub) Put(id string, s *runSession) { + h.mu.Lock() + h.sessions[id] = s + h.mu.Unlock() +} + +func (h *runHub) Get(id string) *runSession { + h.mu.Lock() + defer h.mu.Unlock() + return h.sessions[id] +} + +func (h *runHub) Delete(id string) { + h.mu.Lock() + delete(h.sessions, id) + h.mu.Unlock() +} + +func wantsJSONWait(accept string) bool { + return !strings.Contains(accept, "text/event-stream") +} + +func writeSSE(w http.ResponseWriter, event string, v any) error { + body, err := json.Marshal(v) + if err != nil { + return err + } + if _, err := fmt.Fprintf(w, "event: %s\ndata: %s\n\n", event, body); err != nil { + return err + } + flushSSE(w) + return nil +} + +func flushSSE(w http.ResponseWriter) { + if f, ok := w.(http.Flusher); ok { + f.Flush() + } +} + +func (s *server) handleLogWatch(c *echo.Context) error { + id := c.Param("id") + session := s.runs.Get(id) + if session == nil { + return c.NoContent(http.StatusNotFound) + } + + resp := c.Response() + resp.Header().Set("Content-Type", "text/event-stream") + resp.Header().Set("Cache-Control", "no-cache") + resp.Header().Set("Connection", "keep-alive") + resp.Header().Set("X-Accel-Buffering", "no") + resp.WriteHeader(http.StatusOK) + if _, err := fmt.Fprintf(resp, ": connected\n\n"); err != nil { + return err + } + flushSSE(resp) + + events, cancel := session.Subscribe() + defer cancel() + + ticker := time.NewTicker(15 * time.Second) + defer ticker.Stop() + + ctx := c.Request().Context() + for { + select { + case <-ctx.Done(): + return nil + case <-ticker.C: + if _, err := fmt.Fprintf(resp, ": heartbeat\n\n"); err != nil { + return err + } + flushSSE(resp) + case ev, ok := <-events: + if !ok { + return nil + } + switch ev.Type { + case workertypes.WorkerEventLog: + if err := writeSSE(resp, "log", map[string]string{"line": ev.Line}); err != nil { + return err + } + case workertypes.WorkerEventDone: + if err := writeSSE(resp, "done", ev.Done); err != nil { + return err + } + return nil + case "error": + if err := writeSSE(resp, "done", ev.Done); err != nil { + return err + } + return nil + } + } + } +} + +func applyWorkerEvent(session *runSession, body []byte) { + var evt workertypes.WorkerEvent + if err := json.Unmarshal(body, &evt); err != nil { + log.Printf("could not unmarshal worker event: %v", err) + return + } + switch evt.Type { + case workertypes.WorkerEventLog: + session.AppendLog(evt.Line) + case workertypes.WorkerEventDone, "": + payload := evt.WorkerResponsePayload + if payload == nil { + payload = &workertypes.WorkerResponsePayload{} + if err := json.Unmarshal(body, payload); err != nil { + log.Printf("could not unmarshal done payload: %v", err) + return + } + } + session.Complete(payload) + default: + log.Printf("unknown worker event type %q", evt.Type) + } +} diff --git a/control-service/runs_test.go b/control-service/runs_test.go new file mode 100644 index 0000000..b119b95 --- /dev/null +++ b/control-service/runs_test.go @@ -0,0 +1,141 @@ +package main + +import ( + "io" + "net/http" + "net/http/httptest" + "strings" + "testing" + "time" + + "github.com/mxschmitt/try-playwright/internal/workertypes" +) + +func TestApplyWorkerEventLogAndDone(t *testing.T) { + s := newRunSession("run-1") + applyWorkerEvent(s, []byte(`{"type":"log","line":"hello"}`)) + applyWorkerEvent(s, []byte(`{"type":"done","success":true,"output":"hello","version":"1.2.3","files":[]}`)) + <-s.finished + payload, timedOut := s.Result() + if timedOut { + t.Fatal("timed out") + } + if payload.Output != "hello" || !payload.Success || payload.Version != "1.2.3" { + t.Fatalf("payload %+v", payload) + } +} + +func TestHandleLogWatchSSE(t *testing.T) { + srv := &server{runs: newRunHub()} + srv.initializeHttpServer() + ts := httptest.NewServer(srv.echo) + t.Cleanup(ts.Close) + + session := newRunSession("abc") + srv.runs.Put("abc", session) + + go func() { + time.Sleep(20 * time.Millisecond) + session.AppendLog("early") + session.Complete(&workertypes.WorkerResponsePayload{ + Success: true, + Output: "early", + Version: "1.0.0", + Files: []workertypes.File{}, + }) + }() + + resp, err := http.Get(ts.URL + "/service/control/run/abc/log-watch") + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusOK { + t.Fatalf("status %d", resp.StatusCode) + } + if ct := resp.Header.Get("Content-Type"); !strings.Contains(ct, "text/event-stream") { + t.Fatalf("content-type %q", ct) + } + + body, err := io.ReadAll(resp.Body) + if err != nil { + t.Fatal(err) + } + text := string(body) + if !strings.Contains(text, ": connected") { + t.Fatalf("missing connected comment: %s", text) + } + if !strings.Contains(text, "event: log") || !strings.Contains(text, `"line":"early"`) { + t.Fatalf("missing log event: %s", text) + } + if !strings.Contains(text, "event: done") { + t.Fatalf("missing done event: %s", text) + } +} + +func TestHandleRunAcceptedWithoutJSONWait(t *testing.T) { + workers := &Workers{workers: make(chan *Worker, 1)} + s := &server{ + workers: map[workertypes.WorkerLanguage]*Workers{ + workertypes.WorkerLanguageJavaScript: workers, + }, + runs: newRunHub(), + } + s.initializeHttpServer() + ts := httptest.NewServer(s.echo) + t.Cleanup(ts.Close) + + req, err := http.NewRequest(http.MethodPost, ts.URL+"/service/control/run", strings.NewReader(`{"code":"console.log(1)","language":"javascript"}`)) + if err != nil { + t.Fatal(err) + } + req.Header.Set("Content-Type", "application/json") + req.Header.Set("Accept", "text/event-stream") + resp, err := http.DefaultClient.Do(req) + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusServiceUnavailable { + body, _ := io.ReadAll(resp.Body) + t.Fatalf("status %d want 503 (no worker), body %s", resp.StatusCode, body) + } +} + +func TestWantsJSONWait(t *testing.T) { + if !wantsJSONWait("") || !wantsJSONWait("*/*") || !wantsJSONWait("application/json") { + t.Fatal("expected JSON wait") + } + if wantsJSONWait("text/event-stream") { + t.Fatal("event-stream should not wait") + } +} + +func TestLogWatchUnknownRun(t *testing.T) { + srv := &server{runs: newRunHub()} + srv.initializeHttpServer() + ts := httptest.NewServer(srv.echo) + t.Cleanup(ts.Close) + resp, err := http.Get(ts.URL + "/service/control/run/missing/log-watch") + if err != nil { + t.Fatal(err) + } + defer resp.Body.Close() + if resp.StatusCode != http.StatusNotFound { + t.Fatalf("status %d", resp.StatusCode) + } +} + +func TestApplyWorkerEventLegacyPayload(t *testing.T) { + s := newRunSession("legacy") + applyWorkerEvent(s, []byte(`{"success":true,"output":"2","version":"1.2.3","files":[]}`)) + select { + case <-s.finished: + case <-time.After(time.Second): + t.Fatal("legacy payload did not complete session") + } + payload, _ := s.Result() + if payload.Output != "2" { + t.Fatalf("payload %+v", payload) + } +} diff --git a/control-service/turnstile.go b/control-service/turnstile.go index fa44ac8..7cbbfa0 100644 --- a/control-service/turnstile.go +++ b/control-service/turnstile.go @@ -11,8 +11,8 @@ import ( ) type TurnstileResponse struct { - Success bool `json:"success"` - ErrorCodes []string `json:"error-codes,omitempty"` + Success bool `json:"success"` + ErrorCodes []string `json:"error-codes,omitempty"` } var turnStileHttpClient = &http.Client{Timeout: 15 * time.Second} @@ -26,42 +26,42 @@ func ValidateTurnstile(ctx context.Context, token string, remoteIP string, secre log.Printf("warning: Turnstile remoteIP is empty, skipping validation") return nil } - if token == "" { - return fmt.Errorf("no token provided") - } - requestBody, err := json.Marshal(map[string]string{ - "secret": secretKey, - "response": token, - "remoteip": remoteIP, - }) - if err != nil { - return fmt.Errorf("failed to marshal request body: %w", err) - } + if token == "" { + return fmt.Errorf("no token provided") + } + requestBody, err := json.Marshal(map[string]string{ + "secret": secretKey, + "response": token, + "remoteip": remoteIP, + }) + if err != nil { + return fmt.Errorf("failed to marshal request body: %w", err) + } - req, err := http.NewRequestWithContext(ctx, "POST", "https://challenges.cloudflare.com/turnstile/v0/siteverify", bytes.NewBuffer(requestBody)) - if err != nil { - return fmt.Errorf("failed to create request: %w", err) - } - req.Header.Set("Content-Type", "application/json") + req, err := http.NewRequestWithContext(ctx, "POST", "https://challenges.cloudflare.com/turnstile/v0/siteverify", bytes.NewBuffer(requestBody)) + if err != nil { + return fmt.Errorf("failed to create request: %w", err) + } + req.Header.Set("Content-Type", "application/json") - resp, err := turnStileHttpClient.Do(req) - if err != nil { - return fmt.Errorf("failed to send request: %w", err) - } - defer resp.Body.Close() + resp, err := turnStileHttpClient.Do(req) + if err != nil { + return fmt.Errorf("failed to send request: %w", err) + } + defer resp.Body.Close() - if resp.StatusCode != http.StatusOK { - return fmt.Errorf("unexpected status code: %v", resp.StatusCode) - } + if resp.StatusCode != http.StatusOK { + return fmt.Errorf("unexpected status code: %v", resp.StatusCode) + } - var result TurnstileResponse - if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { - return fmt.Errorf("failed to parse response: %w", err) - } + var result TurnstileResponse + if err := json.NewDecoder(resp.Body).Decode(&result); err != nil { + return fmt.Errorf("failed to parse response: %w", err) + } - if !result.Success { - return fmt.Errorf("turnstile validation failed: %v", result.ErrorCodes) - } + if !result.Success { + return fmt.Errorf("turnstile validation failed: %v", result.ErrorCodes) + } - return nil -} \ No newline at end of file + return nil +} diff --git a/control-service/workers.go b/control-service/workers.go index 0640837..49d7bdb 100644 --- a/control-service/workers.go +++ b/control-service/workers.go @@ -25,7 +25,7 @@ type Workers struct { amqpReplyQueueName string amqpChannel *amqp.Channel k8ClientSet kubernetes.Interface - replies sync.Map // map[string]chan *workertypes.WorkerResponsePayload + replies sync.Map // map[string]*runSession } func newWorkers(language workertypes.WorkerLanguage, workerCount int, k8ClientSet kubernetes.Interface, amqpChannel *amqp.Channel) (*Workers, error) { @@ -78,13 +78,8 @@ func (w *Workers) consumeReplies() error { log.Printf("no reply channel exists for worker %s", msg.CorrelationId) continue } - replyChan := value.(chan *workertypes.WorkerResponsePayload) - var reply *workertypes.WorkerResponsePayload - if err := json.Unmarshal(msg.Body, &reply); err != nil { - log.Printf("could not unmarshal reply json: %v", err) - continue - } - replyChan <- reply + session := value.(*runSession) + applyWorkerEvent(session, msg.Body) } }() return nil @@ -128,9 +123,6 @@ func newWorker(workers *Workers) (*Worker, error) { workers: workers, language: workers.language, } - - w.workers.replies.Store(w.id, make(chan *workertypes.WorkerResponsePayload, 1)) - _, err := w.workers.amqpChannel.QueueDeclare( fmt.Sprintf("rpc_queue_%s", w.id), // name false, // durable @@ -257,14 +249,3 @@ func (w *Worker) Cleanup() error { w.workers.replies.Delete(w.id) return nil } - -func (w *Worker) Subscribe() <-chan *workertypes.WorkerResponsePayload { - value, ok := w.workers.replies.Load(w.id) - if !ok { - // This shouldn't happen, but return a closed channel to avoid panic - ch := make(chan *workertypes.WorkerResponsePayload) - close(ch) - return ch - } - return value.(chan *workertypes.WorkerResponsePayload) -} diff --git a/e2e/tests/api.spec.ts b/e2e/tests/api.spec.ts index 2040b51..e475055 100644 --- a/e2e/tests/api.spec.ts +++ b/e2e/tests/api.spec.ts @@ -191,4 +191,55 @@ class Program expect(body).toHaveProperty('files', []) expect(body).toHaveProperty('output', '2') }) -}) \ No newline at end of file +}) + +test.describe("Live logs", () => { + test("streams stdout over log-watch before done", async ({ request }) => { + const testId = test.info().testId + const code = [ + "console.log('early');", + "const end = Date.now() + 3000;", + "while (Date.now() < end) {}", + "console.log('late');", + ].join('\n') + const started = Date.now() + const startResp = await request.post('/service/control/run', { + headers: { + 'Accept': 'text/event-stream', + 'X-Test-ID': testId, + }, + data: { + code, + language: 'javascript', + }, + timeout: 30 * 1000, + }) + expect(startResp.status()).toBe(202) + const startedBody = await startResp.json() + const id = startedBody.id + expect(id).toBeTruthy() + expect(Date.now() - started).toBeLessThan(2000) + + const watch = await request.get('/service/control/run/' + id + '/log-watch', { + timeout: 30 * 1000, + }) + expect(watch.ok()).toBeTruthy() + const text = await watch.text() + expect(text).toContain(': connected') + expect(text).toContain('event: log') + expect(text).toContain('early') + const doneChunk = text.split('\n\n').find(chunk => chunk.includes('event: done')) + expect(doneChunk).toBeTruthy() + const dataLine = doneChunk.split('\n').find(line => line.startsWith('data: ')) + const donePayload = JSON.parse(dataLine.slice(6)) + expect(donePayload).toMatchObject({ + success: true, + error: '', + files: [], + }) + expect(donePayload.output).toContain('early') + expect(donePayload.output).toContain('late') + expectValidVersion(donePayload) + await attachAggregatorLogs(testId) + }) +}) diff --git a/e2e/tests/visual.spec.ts b/e2e/tests/visual.spec.ts index 6c1b025..dfc4cde 100644 --- a/e2e/tests/visual.spec.ts +++ b/e2e/tests/visual.spec.ts @@ -8,12 +8,9 @@ class TryPlaywrightPage { const panel = this.page.locator('.rs-panel-group > .rs-panel').nth(nth - 1); await panel.getByRole('link').click(); await expect(panel).toHaveClass(/rs-panel-in/); - const responsePromise = this.page.waitForResponse("**/service/control/run"); try { - await Promise.all([ - responsePromise, - this.page.getByRole('button', { name: 'Run' }).click(), - ]); + await this.page.getByRole('button', { name: 'Run' }).click(); + await this.page.getByText(/Duration of \d+ ms with Playwright version|Execution timeout!/).waitFor({ timeout: 120000 }); } finally { await attachAggregatorLogs(); } diff --git a/frontend/src/components/App/index.tsx b/frontend/src/components/App/index.tsx index cb1f999..f162b08 100644 --- a/frontend/src/components/App/index.tsx +++ b/frontend/src/components/App/index.tsx @@ -52,7 +52,12 @@ const App: React.FunctionComponent = () => { const turnstileToken = await gateRef.current!.getToken(turnstileRef.current) setLoading(true) // After await: do not use render-time `code` (stale vs example select). - setResponse(await runCode(getCode(), codeLanguage, turnstileToken)) + const result = await runCode(getCode(), codeLanguage, turnstileToken, (partial) => { + setLoading(false) + onChangeRightPanelMode(false) + setResponse(partial) + }) + setResponse(result) } catch (error) { setResponse({ error: String(error) }) } finally { diff --git a/frontend/src/utils.ts b/frontend/src/utils.ts index bd538e3..2178a32 100644 --- a/frontend/src/utils.ts +++ b/frontend/src/utils.ts @@ -16,11 +16,17 @@ export type ExecutionResponse = Partial<{ output: string; }> -export const runCode = async (code: string, codeLanguage: CodeLanguage, turnstileToken: string): Promise => { +export const runCode = async ( + code: string, + codeLanguage: CodeLanguage, + turnstileToken: string, + onUpdate?: (resp: ExecutionResponse) => void, +): Promise => { if (codeLanguage === CodeLanguage.PLAYWRIGHT_TEST) codeLanguage = CodeLanguage.JAVASCRIPT const headers: Record = { "Content-Type": "application/json", + "Accept": "text/event-stream", } if (window.__TRY_PLAYWRIGHT_TEST_ID__) { headers["X-Test-ID"] = window.__TRY_PLAYWRIGHT_TEST_ID__ @@ -35,10 +41,17 @@ export const runCode = async (code: string, codeLanguage: CodeLanguage, turnstil }) }) - if (!resp.ok) { - if (resp.status === 429) { - return { error: "You are rate limited, please try again in a few minutes." } + if (resp.status === 429) { + return { error: "You are rate limited, please try again in a few minutes." } + } + if (resp.status === 202) { + const body = await resp.json() as { id?: string } + if (!body.id) { + return { error: "Execution was not successful, please try again in a few minutes." } } + return watchRunLogs(body.id, onUpdate) + } + if (!resp.ok) { if (resp.headers.get("Content-Type")?.includes("application/json")) { return await resp.json() } @@ -47,6 +60,43 @@ export const runCode = async (code: string, codeLanguage: CodeLanguage, turnstil return await resp.json() } +const watchRunLogs = (id: string, onUpdate?: (resp: ExecutionResponse) => void): Promise => { + return new Promise((resolve, reject) => { + const es = new EventSource(`/service/control/run/${encodeURIComponent(id)}/log-watch`) + let output = "" + let settled = false + const finish = (value: ExecutionResponse) => { + if (settled) { + return + } + settled = true + es.close() + resolve(value) + } + es.addEventListener("log", (ev) => { + const data = JSON.parse((ev as MessageEvent).data) as { line?: string } + const line = data.line ?? "" + output = output ? `${output}\n${line}` : line + onUpdate?.({ output }) + }) + es.addEventListener("done", (ev) => { + const data = JSON.parse((ev as MessageEvent).data) as ExecutionResponse + if (data.output === undefined || data.output === "") { + data.output = output + } + finish(data) + }) + es.onerror = () => { + if (settled) { + return + } + settled = true + es.close() + reject(new Error("log watch failed")) + } + }) +} + declare global { interface Window { gtag?: (kind: string, event: string, metaData: Record) => void diff --git a/internal/worker/worker.go b/internal/worker/worker.go index cbd17f7..23283fd 100644 --- a/internal/worker/worker.go +++ b/internal/worker/worker.go @@ -1,6 +1,7 @@ package worker import ( + "bufio" "bytes" "encoding/json" "errors" @@ -12,6 +13,7 @@ import ( "os/exec" "path/filepath" "strings" + "sync" "time" "github.com/mxschmitt/try-playwright/internal/logagg" @@ -25,6 +27,9 @@ type executionHandler func(worker *Worker, code string) error type Worker struct { options *WorkerExecutionOptions channel *amqp.Channel + pubMu sync.Mutex + replyTo string + corrID string TmpDir string requestID string testID string @@ -32,6 +37,7 @@ type Worker struct { output *bytes.Buffer files []string env []string + onLog func(line string) } var ( @@ -149,15 +155,39 @@ func (w *Worker) ExecCommand(name string, args ...string) error { env = append(env, e...) } + env = append(env, "PYTHONUNBUFFERED=1") + + cmdPath := path + cmdArgs := append([]string{name}, args...) + if stdbuf, err := exec.LookPath("stdbuf"); err == nil { + cmdPath = stdbuf + cmdArgs = append([]string{"stdbuf", "-oL", "-eL", path}, args...) + } + + pr, pw := io.Pipe() c := exec.Cmd{ Dir: w.TmpDir, - Path: path, - Args: append([]string{name}, args...), - Stdout: io.MultiWriter(os.Stdout, w.output), - Stderr: io.MultiWriter(os.Stderr, w.output), + Path: cmdPath, + Args: cmdArgs, + Stdout: io.MultiWriter(os.Stdout, pw), + Stderr: io.MultiWriter(os.Stderr, pw), Env: env, } - if err := c.Run(); err != nil { + + scanDone := make(chan struct{}) + go func() { + defer close(scanDone) + scanner := bufio.NewScanner(pr) + scanner.Buffer(make([]byte, 0, 64*1024), 1024*1024) + for scanner.Scan() { + w.emitLog(scanner.Text()) + } + }() + + runErr := c.Run() + _ = pw.Close() + <-scanDone + if runErr != nil { return errors.New("could not run command") } files, err := collector.Collect() @@ -179,6 +209,8 @@ func (w *Worker) consumeMessage(incomingMessages <-chan amqp.Delivery) error { } w.requestID = incomingMessageParsed.RequestID w.testID = incomingMessageParsed.TestID + w.replyTo = incomingMessage.ReplyTo + w.corrID = incomingMessage.CorrelationId if w.requestID != "" { w.AddEnv("PLAYWRIGHT_REQUEST_ID", w.requestID) } @@ -210,21 +242,10 @@ func (w *Worker) consumeMessage(incomingMessages <-chan amqp.Delivery) error { outgoingMessage.Output = w.options.TransformOutput(w.output.String()) outgoingMessage.RequestID = w.requestID outgoingMessage.TestID = w.testID - outgoingMessageBody, err := json.Marshal(outgoingMessage) - if err != nil { - return fmt.Errorf("could not marshal outgoing message payload: %w", err) - } - err = w.channel.Publish( - "", // exchange - incomingMessage.ReplyTo, // routing key - false, // mandatory - false, // immediate - amqp.Publishing{ - ContentType: "application/json", - CorrelationId: incomingMessage.CorrelationId, - Body: outgoingMessageBody, - }) - if err != nil { + if err := w.publishEvent(workertypes.WorkerEvent{ + Type: workertypes.WorkerEventDone, + WorkerResponsePayload: outgoingMessage, + }); err != nil { return fmt.Errorf("could not publish message: %w", err) } @@ -234,6 +255,45 @@ func (w *Worker) consumeMessage(incomingMessages <-chan amqp.Delivery) error { return nil } +func (w *Worker) emitLog(line string) { + if w.onLog != nil { + w.onLog(line) + } + w.output.WriteString(line) + w.output.WriteByte('\n') + if err := w.publishEvent(workertypes.WorkerEvent{ + Type: workertypes.WorkerEventLog, + Line: line, + }); err != nil && w.logger != nil { + w.logger.WithError(err).Warn("could not publish log event") + } +} + +func (w *Worker) publishEvent(evt workertypes.WorkerEvent) error { + if w.channel == nil || w.replyTo == "" { + return nil + } + body, err := json.Marshal(evt) + if err != nil { + return fmt.Errorf("could not marshal event: %w", err) + } + w.pubMu.Lock() + defer w.pubMu.Unlock() + if err := w.channel.Publish( + "", + w.replyTo, + false, + false, + amqp.Publishing{ + ContentType: "application/json", + CorrelationId: w.corrID, + Body: body, + }); err != nil { + return fmt.Errorf("could not publish event: %w", err) + } + return nil +} + var uploadFilesEndpoint = fmt.Sprintf("%s/api/v1/file/upload", os.Getenv("FILE_SERVICE_URL")) func (w *Worker) uploadFiles() ([]workertypes.File, error) { diff --git a/internal/worker/worker_test.go b/internal/worker/worker_test.go index 96e0211..e05d217 100644 --- a/internal/worker/worker_test.go +++ b/internal/worker/worker_test.go @@ -2,7 +2,9 @@ package worker import ( "errors" + "sync" "testing" + "time" amqp "github.com/rabbitmq/amqp091-go" ) @@ -23,3 +25,38 @@ func TestConsumeMessageClosedDeliveryChannel(t *testing.T) { t.Fatalf("expected errAMQPChannelClosed, got %v", err) } } + +func TestExecCommandStreamsLogsBeforeExit(t *testing.T) { + w := NewWorker(&WorkerExecutionOptions{ + Handler: func(worker *Worker, code string) error { return nil }, + }) + w.TmpDir = t.TempDir() + + early := make(chan struct{}) + var once sync.Once + w.onLog = func(line string) { + if line == "early" { + once.Do(func() { close(early) }) + } + } + + done := make(chan error, 1) + go func() { + done <- w.ExecCommand("sh", "-c", "echo early; sleep 1; echo late") + }() + + select { + case <-early: + case err := <-done: + t.Fatalf("command finished before early log: %v", err) + case <-time.After(3 * time.Second): + t.Fatal("did not receive streamed log before command exit") + } + + if err := <-done; err != nil { + t.Fatalf("ExecCommand: %v", err) + } + if got := w.options.TransformOutput(w.output.String()); got != "early\nlate" { + t.Fatalf("output %q", got) + } +} diff --git a/internal/workertypes/types.go b/internal/workertypes/types.go index 90e8552..8c0fba0 100644 --- a/internal/workertypes/types.go +++ b/internal/workertypes/types.go @@ -2,6 +2,11 @@ package workertypes import "slices" +const ( + WorkerEventLog = "log" + WorkerEventDone = "done" +) + type File struct { PublicURL string `json:"publicURL"` FileName string `json:"fileName"` @@ -19,6 +24,15 @@ type WorkerResponsePayload struct { TestID string `json:"testId"` } +// WorkerEvent is published on the AMQP reply queue while a job runs. +// Type is "log" (Line set) or "done" (embedded payload). An empty Type is +// treated as done so older workers still complete a run. +type WorkerEvent struct { + Type string `json:"type"` + Line string `json:"line,omitempty"` + *WorkerResponsePayload +} + type WorkerRequestPayload struct { Token string `json:"token"` Code string `json:"code"` diff --git a/worker-python/main.go b/worker-python/main.go index 1022e5e..d5ca83e 100644 --- a/worker-python/main.go +++ b/worker-python/main.go @@ -5,7 +5,7 @@ import ( ) func handler(w *worker.Worker, code string) error { - return w.ExecCommand("python", "-c", code) + return w.ExecCommand("python", "-u", "-c", code) } func main() {