diff --git a/FORK.md b/FORK.md index a62bd550fd..93e002d8fd 100644 --- a/FORK.md +++ b/FORK.md @@ -3,11 +3,17 @@ Base: upstream tag `release-1.1.1` (commit `531700ca`), the stable 1.1.1 umbrella release. -## Patch carried here +## Patches carried here -**execd ends a Jupyter run when its kernel stops.** A run whose kernel stops -before replying never ended. Deleting the context of a running cell makes the -kernel broadcast `shutdown_reply` and an `idle` status, and then it is killed. +Two patches to how execd runs code in a Jupyter kernel +(`components/execd/pkg/jupyter/execute`). Both make a run's `/code` stream end +when that run's execution ends, and no sooner. + +### 1. execd ends a Jupyter run when its kernel stops + +A run whose kernel stops before replying never ended. Deleting the context of +a running cell makes the kernel broadcast `shutdown_reply` and an `idle` +status, and then it is killed. It never sends the cell's `execute_reply`, and Jupyter Server keeps the kernel websocket open, so `finalizeExecution` waited for that reply forever and the `/code` stream stayed open, sending only pings, until the client gave up. A @@ -27,7 +33,9 @@ run that had already replied is finalized without the extra error. As in Where the stream used to hang, a client now sees the events of an ordinary failing cell: `status`, `error` (`ename` `KernelStopped`), then -`execution_complete`, and the stream closes. Nothing else changes. +`execution_complete`, and the stream closes. Nothing else changes. (Before +patch 2, an `idle` for another request could start finalization first, and +`execution_complete` then came before the error.) What this does not cover: @@ -40,13 +48,9 @@ What this does not cover: - **A websocket that drops mid-run.** `receiveMessages` stops on a read error and the run waits until the client gives up, as before. Ending the run there needs the reader to know which run owns the connection (a late read error - from one run's closed socket must not end the next run on the same client), - so it is left out of this patch. -- **Ordering after a stray `idle`.** Handlers do not match status messages to - the request's `msg_id` (unchanged upstream behaviour). If an unrelated - `idle` starts finalization before the kernel stops, `execution_complete` - can arrive before the `KernelStopped` error; the error still arrives before - the stream closes. + from one run's closed socket must not end the next run on the same client). + Since patch 2 each reader reads only its own connection, but a read error + still ends nothing. The `components/execd/pkg/jupyter/execute` package is unchanged between `release-1.1.1-rc.1`, this tag, upstream `main` at the time of the cut, and @@ -54,6 +58,114 @@ The `components/execd/pkg/jupyter/execute` package is unchanged between Drop this patch when an upstream release ends a run whose kernel stops. +### 2. execd matches messages to the run's request and wakes a stalled kernel + +About one `/code` run in 50 on an existing context never ended. execd +reported `execution_complete` almost at once, then the stream sent only +pings. The context stayed busy for good: interrupting did nothing, and +deleting the context ended the stream only through patch 1. Two defects +combine here. + +**The kernel left the request unread.** The kernel never read the run's +`execute_request`: it published no `busy`, no output and no reply, and Jupyter +Server reported it idle. The request was still queued on the kernel's shell +socket; opening any new websocket to the kernel made the kernel read it, and +the run then finished with the right result. ipykernel 7.0 to 7.3 reads shell +requests through a `ZMQStream` but sends replies directly on the same socket, +and that send can consume the read event for a request arriving at the same +moment ([ipython/ipykernel#1529](https://github.com/ipython/ipykernel/pull/1529), +fixed in ipykernel 7.4.0). The request is then not read until something else +wakes the socket. `opensandbox/code-interpreter:v1.1.0` ships ipykernel 7.3.0. +execd makes the collision common: it connects to the kernel once per run; +on every new websocket Jupyter Server sends a `kernel_info_request` on the +shell and control channels (unless it believes the kernel is busy); and execd +sends the run's request as soon as it is connected, so the shell reply to that +`kernel_info_request` often goes out just as the request arrives. Re-arming +the kernel's stream after each reply (a local patch to the kernel, to the same +effect as the ipykernel 7.4.0 fix) removed every hang (0 in 300 runs). + +If nothing for the run's request has arrived 500 ms after it was sent, execd now +opens and closes another websocket to the kernel; Jupyter Server connects new +sockets to the kernel for it, which wakes the shell socket. It tries up to four +times, doubling the delay (after 0.5, 1.5, 3.5 and 7.5 s). Before each attempt +it stops if anything for the request has arrived, if the run has ended (its +result channel is closed), or if the run's connection has been closed, as +happens when the run's context is cancelled. A kernel that picked up the +request is never woken, however long the run takes. Sending another message on +the run's own connection does not work (tried: every stalled request stayed +unread through all four attempts): it queues behind the unread request on the +same zmq pipe. + +The 500 ms count from when execd writes the request, and a run can receive +nothing for longer without any stall: Jupyter Server forwards the request only +after its own `kernel_info_request` round on the new websocket, and a kernel +still busy with an earlier request (a run whose client gave up, say) sends +nothing for the new one until it finishes. Such a run gets up to four +wake-ups, each a websocket opened and closed (Jupyter Server sends no +`kernel_info_request` to a kernel it believes busy), and an info-level log line +that says nothing has arrived yet. + +**Messages were not matched to the request.** A run's websocket also carries +messages for other requests: the `busy` and `idle` statuses for Jupyter +Server's `kernel_info_request`, usually before the run's own, and the output +and reply of an earlier run that is still going. The handlers ignored +`parent_header`, so that `idle` started `finalizeExecution` on almost every +run and `execution_complete` could come before the result (in 24 of the 300 +upstream runs below), and another request's output, error or execution count +was reported as the run's. An earlier request's reply could also end +finalization before the run's own result arrived, and an error in it made a +later kernel stop look like a run that had already replied, so no +`KernelStopped` error was sent. Replies, output and `idle` statuses now count +only when their parent is the run's `execute_request`. `shutdown_reply`, +status `restarting` and status `dead` still end the run whatever the parent +(patch 1). + +Each kernel websocket is also now read only by the reader started for it. +The reader used to look up the client's current connection before every +read, so a reader still handling a message when its run disconnected could +go on reading the next run's connection alongside that connection's own +reader. + +Measured with a local `opensandbox-server` 1.1.1 and +`opensandbox/code-interpreter:v1.1.0` (ipykernel 7.3.0), running `z = i; z` +back to back on one context through the Python SDK 1.1.1, a run counted as hung +when its stream had not ended after 8 s: + +| execd | runs | hung | `execution_complete` before the result | +| --- | --- | --- | --- | +| upstream `release-1.1.1` | 300 | 6 (deleting the context did not end them) | 24 | +| patch 1 only | 300 | 6 (deleting the context ended them, `KernelStopped`) | 43 | +| patches 1 and 2 | 600 | 0 | 0 | +| patches 1 and 2, after review | 600 | 0 | 0 | + +With both patches the kernel left 10 and 19 of the requests unread in the two +runs; one wake-up each made it read them, those runs took about 0.5 s longer, +and every run returned its own result. + +What this does not cover: + +- **A kernel that stays stuck.** After four wake-ups (7.5 s) execd stops + trying, and the run waits as before. Every stalled request seen needed one. +- **An aborted request.** A kernel that aborts a request replies with status + `aborted` and no execution count, and `finalizeExecution` keeps waiting + (unchanged upstream behaviour). ipykernel aborts only requests queued behind + one that failed, so this needs a run sent while an earlier run on the same + context is still executing. +- **Kernel-wide events between connecting and registering.** A `dead` or + `restarting` status that reaches a run between its connect and its handler + registration (microseconds) goes to the previous run's handlers and is + lost to this run. + +`components/execd/pkg/jupyter/execute` and `pkg/runtime/jupyter.go` are +unchanged between this tag and upstream `main` on 2026-10-20. + +Drop the message matching and the reader change when an upstream release +matches a run's messages to its request. Drop the wake-up when the kernels +execd drives run ipykernel 7.4.0 or later. `opensandbox/code-interpreter:v1.1.0` +can select Python 3.10 to 3.14, each with ipykernel 7.3.0, and ipykernel 7.4.0 +requires Python 3.11 or later, so that cannot happen while the image keeps its +Python 3.10 environment. + ## Rules for this branch - Do not merge or fast-forward upstream `main` into it. Re-cut it onto a diff --git a/components/execd/pkg/jupyter/execute/execute.go b/components/execd/pkg/jupyter/execute/execute.go index bf67c9e481..50cb087887 100644 --- a/components/execd/pkg/jupyter/execute/execute.go +++ b/components/execd/pkg/jupyter/execute/execute.go @@ -28,6 +28,27 @@ import ( "github.com/gorilla/websocket" execdflag "github.com/alibaba/opensandbox/execd/pkg/flag" + "github.com/alibaba/opensandbox/execd/pkg/log" +) + +// A kernel can leave a request unread on its shell socket. ipykernel 7.0 to 7.3 sends replies +// directly on the socket it reads requests from, and that send can consume the read event for a +// request arriving at the same moment (ipython/ipykernel#1529, fixed in 7.4.0). The request then +// stays unread until something else wakes the socket, such as a new connection to it. On every new +// websocket Jupyter Server sends a kernel_info_request on the shell and control channels (unless it +// believes the kernel is busy), and a run connects just before it sends its request, so the shell +// reply to that kernel_info_request regularly meets the request. +// +// A run that has received nothing for its request kernelWakeDelay after sending it opens and closes +// another websocket: Jupyter Server connects new sockets to the kernel for it, which wakes the +// shell socket. Another message on the run's own connection would not: it arrives on the same zmq +// pipe as the unread request, which raises no new read event. The delay doubles after each of +// kernelWakeAttempts attempts. Waking stops when anything for the request arrives, when the run +// ends, or when the run's connection is closed. A kernel still busy with an earlier request also +// sends nothing for a while; waking it costs a websocket and is otherwise harmless. +var ( + kernelWakeDelay = 500 * time.Millisecond + kernelWakeAttempts = 4 ) type HTTPClient interface { @@ -38,6 +59,8 @@ type Client struct { httpClient HTTPClient conn *websocket.Conn + // connClosed is closed when conn is disconnected. + connClosed chan struct{} handlers map[MessageType]func(*Message) @@ -73,10 +96,11 @@ func (c *Client) Connect(wsURL string) error { return fmt.Errorf("failed to connect to kernel: %w", err) } c.conn = conn + c.connClosed = make(chan struct{}) c.registerDefaultHandlers() - safego.Go(func() { c.receiveMessages() }) + safego.Go(func() { c.receiveMessages(conn) }) return nil } @@ -88,6 +112,7 @@ func (c *Client) Disconnect() { if c.conn != nil { c.conn.Close() c.conn = nil + close(c.connClosed) } } @@ -104,6 +129,17 @@ type streamExecutionState struct { executeMutex sync.Mutex resultMutex sync.Mutex + // msgID is the msg_id of this run's execute_request. The kernel sets it as the parent of every + // message that answers the request. The run's connection also carries messages for other + // requests: the busy and idle statuses for Jupyter Server's kernel_info_request on connect, and + // the output of an earlier run that is still going. Those are not this run's. + msgID string + // pickedUp is closed by the first message for msgID: the kernel is working on the request. + pickedUp chan struct{} + pickedUpOnce sync.Once + // ended is closed with resultChan. + ended chan struct{} + // chanMutex guards resultChan's close-vs-send race. state.result.ExecutionCount/Error can be // set by handleExecuteReply with no corresponding resultChan send at all (e.g. plain print() // cells never emit execute_result), so finalizeExecution's poll loop can observe "done" and @@ -124,9 +160,20 @@ func newStreamExecutionState(startTime time.Time) *streamExecutionState { Stream: make([]*StreamOutput, 0), ExecutionTime: 0, }, + pickedUp: make(chan struct{}), + ended: make(chan struct{}), } } +// owns reports whether msg answers this run's request, and records that the kernel picked it up. +func (state *streamExecutionState) owns(msg *Message) bool { + if msg.ParentHeader.MessageID != state.msgID { + return false + } + state.pickedUpOnce.Do(func() { close(state.pickedUp) }) + return true +} + // trySend delivers notify on resultChan unless the channel has already been closed. It must be // used for every send on resultChan so sends can never race closeResultChan. func (state *streamExecutionState) trySend(resultChan chan *ExecutionResult, notify *ExecutionResult) { @@ -146,6 +193,7 @@ func (state *streamExecutionState) closeResultChan(resultChan chan *ExecutionRes defer state.chanMutex.Unlock() state.closed = true close(resultChan) + close(state.ended) } func (c *Client) ExecuteCodeStream(code string, resultChan chan *ExecutionResult) error { @@ -159,6 +207,7 @@ func (c *Client) ExecuteCodeStream(code string, resultChan chan *ExecutionResult } state := newStreamExecutionState(time.Now()) + state.msgID = msg.Header.MessageID c.clearTemporaryHandlers() c.registerExecuteCodeStreamHandlers(state, resultChan) @@ -167,9 +216,70 @@ func (c *Client) ExecuteCodeStream(code string, resultChan chan *ExecutionResult return fmt.Errorf("failed to send execution request: %w", err) } + c.mu.Lock() + connClosed := c.connClosed + c.mu.Unlock() + delay, attempts := kernelWakeDelay, kernelWakeAttempts + safego.Go(func() { c.wakeKernelUntilPickedUp(state, connClosed, delay, attempts) }) + return nil } +// wakeKernelUntilPickedUp wakes a kernel that may have left the run's request unread, see +// kernelWakeDelay. connClosed is closed when the run's connection is. +func (c *Client) wakeKernelUntilPickedUp(state *streamExecutionState, connClosed <-chan struct{}, delay time.Duration, attempts int) { + for attempt := 1; attempt <= attempts; attempt++ { + timer := time.NewTimer(delay) + select { + case <-state.pickedUp: + case <-state.ended: + case <-connClosed: + case <-timer.C: + } + timer.Stop() + if state.wakeNotNeeded(connClosed) { + return + } + log.Info("jupyter: nothing yet for request %s after %s; opening another kernel connection in case the kernel left it unread (attempt %d of %d)", + state.msgID, time.Since(state.startTime).Round(time.Millisecond), attempt, attempts) + c.wakeKernel() + delay *= 2 + } +} + +// wakeNotNeeded reports whether the kernel picked up the request, the run ended, or the run's +// connection was closed. +func (state *streamExecutionState) wakeNotNeeded(connClosed <-chan struct{}) bool { + select { + case <-state.pickedUp: + return true + case <-state.ended: + return true + case <-connClosed: + return true + default: + return false + } +} + +// wakeKernel opens and closes a websocket to the kernel. Jupyter Server connects new sockets to the +// kernel for it, and the kernel's shell socket then reads any request it left unread. +func (c *Client) wakeKernel() { + c.mu.Lock() + wsURL := c.wsURL + c.mu.Unlock() + + conn, resp, err := websocket.DefaultDialer.Dial(wsURL, nil) + if resp != nil && err != nil { + resp.Body.Close() + } + if err != nil { + log.Warn("jupyter: wake kernel: %v", err) + return + } + conn.Close() +} + func (c *Client) buildExecuteMessage(code string) (*Message, error) { msgID := c.nextMessageID() request := &ExecuteRequest{ @@ -204,18 +314,31 @@ func (c *Client) buildExecuteMessage(code string) (*Message, error) { return msg, nil } +// registerExecuteCodeStreamHandlers registers the run's handlers. They are in place before the +// request is sent, so nothing the kernel sends for it can arrive first. Replies, output and idle +// statuses count only when they answer the run's request; a stopped kernel ends the run whatever +// the parent, since its shutdown_reply answers Jupyter Server's shutdown_request and its restarting +// and dead statuses come from Jupyter Server itself. func (c *Client) registerExecuteCodeStreamHandlers(state *streamExecutionState, resultChan chan *ExecutionResult) { c.registerHandler(MsgExecuteReply, func(msg *Message) { - c.handleExecuteReply(msg, state) + if state.owns(msg) { + c.handleExecuteReply(msg, state) + } }) c.registerHandler(MsgExecuteResult, func(msg *Message) { - c.handleExecuteResult(msg, state, resultChan) + if state.owns(msg) { + c.handleExecuteResult(msg, state, resultChan) + } }) c.registerHandler(MsgStream, func(msg *Message) { - c.handleStreamOutput(msg, state, resultChan) + if state.owns(msg) { + c.handleStreamOutput(msg, state, resultChan) + } }) c.registerHandler(MsgError, func(msg *Message) { - c.handleExecutionError(msg, state, resultChan) + if state.owns(msg) { + c.handleExecutionError(msg, state, resultChan) + } }) c.registerHandler(MsgStatus, func(msg *Message) { c.handleExecutionStatus(msg, state, resultChan) @@ -311,8 +434,9 @@ func (c *Client) handleExecutionStatus(msg *Message, state *streamExecutionState case StateDead: c.handleKernelStopped(state, resultChan, "the kernel died before the execution finished") return - case StateIdle: - default: + } + // Any other status answers one request, and only the idle for this run's request ends it. + if !state.owns(msg) || status.ExecutionState != StateIdle { return } @@ -506,21 +630,15 @@ func (c *Client) clearTemporaryHandlers() { c.registerDefaultHandlers() } -func (c *Client) receiveMessages() { +// receiveMessages dispatches the messages read from conn until it is closed. It reads only the +// connection it was started for: a reader still handling a message when its run disconnects must +// not go on to read the next run's connection alongside that connection's own reader. +func (c *Client) receiveMessages(conn *websocket.Conn) { for { - c.mu.Lock() - conn := c.conn - c.mu.Unlock() - - if conn == nil { - break - } - var msg Message - err := conn.ReadJSON(&msg) - if err != nil { + if err := conn.ReadJSON(&msg); err != nil { // connection may already be closed - break + return } c.handleMessage(&msg) diff --git a/components/execd/pkg/jupyter/execute/execute_test.go b/components/execd/pkg/jupyter/execute/execute_test.go index fc7570d030..59b2b25640 100644 --- a/components/execd/pkg/jupyter/execute/execute_test.go +++ b/components/execd/pkg/jupyter/execute/execute_test.go @@ -19,6 +19,8 @@ import ( "net/http" "net/http/httptest" "strings" + "sync" + "sync/atomic" "testing" "time" @@ -346,13 +348,13 @@ func TestExecuteCodeStreamEndsWhenKernelStopsMidExecution(t *testing.T) { return } busy, _ := json.Marshal(StatusUpdate{ExecutionState: StateBusy}) - require.NoError(t, conn.WriteJSON(Message{ + busyMsg := Message{ Header: Header{MessageID: "busy-id", Session: executeRequest.Header.Session, MessageType: string(MsgStatus)}, ParentHeader: executeRequest.Header, Content: busy, - })) - for _, msg := range tc.messages(executeRequest.Header.Session, executeRequest.Header) { - require.NoError(t, conn.WriteJSON(msg)) + } + if !writeMessages(t, conn, append([]Message{busyMsg}, tc.messages(executeRequest.Header.Session, executeRequest.Header)...)...) { + return } // Jupyter Server leaves the websocket open after the kernel is gone. <-release @@ -393,10 +395,10 @@ func TestExecuteCodeStreamEndsWhenKernelStopsMidExecution(t *testing.T) { } } -// TestExecuteCodeStreamEndsWhenKernelStopsWhileFinalizing covers a kernel that stops after an idle -// status has already started finalizeExecution, so finalization is polling for a reply that will -// never come. The KernelStopped error must still reach the client before the stream closes: the -// error is sent before it is recorded, and recording it is what ends the poll. +// TestExecuteCodeStreamEndsWhenKernelStopsWhileFinalizing covers a kernel that stops after the +// run's idle status has already started finalizeExecution, so finalization is polling for a reply +// that will never come. The KernelStopped error must still reach the client before the stream +// closes: the error is sent before it is recorded, and recording it is what ends the poll. func TestExecuteCodeStreamEndsWhenKernelStopsWhileFinalizing(t *testing.T) { previousPollInterval := execdflag.JupyterIdlePollInterval execdflag.JupyterIdlePollInterval = 5 * time.Millisecond @@ -413,17 +415,22 @@ func TestExecuteCodeStreamEndsWhenKernelStopsWhileFinalizing(t *testing.T) { } session := executeRequest.Header.Session idle, _ := json.Marshal(StatusUpdate{ExecutionState: StateIdle}) - require.NoError(t, conn.WriteJSON(Message{ - Header: Header{MessageID: "stray-idle-id", Session: session, MessageType: string(MsgStatus)}, - Content: idle, - })) + if !writeMessages(t, conn, Message{ + Header: Header{MessageID: "idle-id", Session: session, MessageType: string(MsgStatus)}, + ParentHeader: executeRequest.Header, + Content: idle, + }) { + return + } // Several poll intervals, so finalizeExecution is polling when the kernel stops. time.Sleep(50 * time.Millisecond) restarting, _ := json.Marshal(StatusUpdate{ExecutionState: StateRestarting}) - require.NoError(t, conn.WriteJSON(Message{ + if !writeMessages(t, conn, Message{ Header: Header{MessageID: "restarting-id", Session: session, MessageType: string(MsgStatus)}, Content: restarting, - })) + }) { + return + } <-release }) defer server.Close() @@ -455,3 +462,541 @@ func TestExecuteCodeStreamEndsWhenKernelStopsWhileFinalizing(t *testing.T) { } } } + +// kernelMessage builds a message as the kernel would send it. Test servers call it from their +// handler goroutines, so it reports a failure with t.Errorf rather than stopping the test. +func kernelMessage(t *testing.T, msgType MessageType, id, session string, parent Header, content interface{}) Message { + t.Helper() + raw, err := json.Marshal(content) + if err != nil { + t.Errorf("failed to marshal %s content: %v", msgType, err) + } + return Message{ + Header: Header{MessageID: id, Session: session, MessageType: string(msgType)}, + ParentHeader: parent, + Content: raw, + } +} + +// writeMessages sends msgs from a test server's handler goroutine, which must not call t.FailNow. +// It reports whether every write succeeded, so the handler can return on failure. +func writeMessages(t *testing.T, conn *websocket.Conn, msgs ...Message) bool { + t.Helper() + for _, msg := range msgs { + if err := conn.WriteJSON(msg); err != nil { + t.Errorf("failed to write %s: %v", msg.Header.MessageType, err) + return false + } + } + return true +} + +// collectResults reads resultChan until it is closed, failing the test if that takes too long. +func collectResults(t *testing.T, resultChan chan *ExecutionResult) []*ExecutionResult { + t.Helper() + var results []*ExecutionResult + deadline := time.After(2 * time.Second) + for { + select { + case result, open := <-resultChan: + if !open { + return results + } + results = append(results, result) + case <-deadline: + t.Fatal("the stream did not end") + } + } +} + +// TestExecuteCodeStreamIgnoresIdleOfAnotherRequest covers the idle statuses that reach a run before +// its own: Jupyter Server sends a kernel_info_request on every new websocket and the kernel's busy +// and idle for it arrive on the run's connection. Such an idle must not start finalization, which +// would report the run complete before its result. The run ends on its own idle and reply. +func TestExecuteCodeStreamIgnoresIdleOfAnotherRequest(t *testing.T) { + previousPollInterval := execdflag.JupyterIdlePollInterval + execdflag.JupyterIdlePollInterval = time.Millisecond + t.Cleanup(func() { + execdflag.JupyterIdlePollInterval = previousPollInterval + }) + + release := make(chan struct{}) + server := createTestServer(t, func(conn *websocket.Conn) { + var executeRequest Message + if err := conn.ReadJSON(&executeRequest); err != nil { + t.Errorf("failed to read execution request: %v", err) + return + } + session := executeRequest.Header.Session + own := executeRequest.Header + kernelInfo := Header{MessageID: "kernel-info-request-id", MessageType: string(MsgKernelInfo)} + if !writeMessages(t, conn, + kernelMessage(t, MsgStatus, "info-busy", session, kernelInfo, StatusUpdate{ExecutionState: StateBusy}), + kernelMessage(t, MsgStatus, "info-idle", session, kernelInfo, StatusUpdate{ExecutionState: StateIdle}), + ) { + return + } + // Many poll intervals: a run that took that idle as its own would end here. + time.Sleep(50 * time.Millisecond) + if !writeMessages(t, conn, + kernelMessage(t, MsgStatus, "own-busy", session, own, StatusUpdate{ExecutionState: StateBusy}), + kernelMessage(t, MsgExecuteResult, "own-result", session, own, ExecuteResult{ + ExecutionCount: 2, + Data: map[string]interface{}{"text/plain": "own"}, + }), + kernelMessage(t, MsgExecuteReply, "own-reply", session, own, ExecuteReply{ExecutionCount: 2, Status: "ok"}), + kernelMessage(t, MsgStatus, "own-idle", session, own, StatusUpdate{ExecutionState: StateIdle}), + ) { + return + } + <-release + }) + defer server.Close() + defer close(release) + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/kernels/test-kernel-id/channels" + client := NewClient("", nil) + require.NoError(t, client.Connect(wsURL)) + defer client.Disconnect() + + resultChan := make(chan *ExecutionResult, 10) + require.NoError(t, client.ExecuteCodeStream("x", resultChan)) + + results := collectResults(t, resultChan) + resultAt, completeAt := -1, -1 + for i, result := range results { + if result.ExecutionData != nil && resultAt < 0 { + resultAt = i + } + if result.ExecutionTime > 0 && completeAt < 0 { + completeAt = i + } + } + require.GreaterOrEqual(t, resultAt, 0, "expected the run's own result") + require.Equal(t, "own", results[resultAt].ExecutionData["text/plain"]) + require.Greater(t, completeAt, resultAt, "the run must not complete before its own result") +} + +// TestExecuteCodeStreamIgnoresOutputOfAnotherRequest covers output that belongs to an earlier +// request: an earlier run that the client gave up on can still be running and finish while this +// run waits. Its output and reply must not be reported as this run's, nor end it. +func TestExecuteCodeStreamIgnoresOutputOfAnotherRequest(t *testing.T) { + previousPollInterval := execdflag.JupyterIdlePollInterval + execdflag.JupyterIdlePollInterval = time.Millisecond + t.Cleanup(func() { + execdflag.JupyterIdlePollInterval = previousPollInterval + }) + + release := make(chan struct{}) + server := createTestServer(t, func(conn *websocket.Conn) { + var executeRequest Message + if err := conn.ReadJSON(&executeRequest); err != nil { + t.Errorf("failed to read execution request: %v", err) + return + } + session := executeRequest.Header.Session + own := executeRequest.Header + earlier := Header{MessageID: "earlier-execute-request-id", MessageType: string(MsgExecuteRequest)} + if !writeMessages(t, conn, + kernelMessage(t, MsgStream, "earlier-stream", session, earlier, StreamOutput{Name: StreamStdout, Text: "earlier\n"}), + kernelMessage(t, MsgError, "earlier-error", session, earlier, ErrorOutput{EName: "EarlierError", EValue: "earlier"}), + kernelMessage(t, MsgExecuteResult, "earlier-result", session, earlier, ExecuteResult{ + ExecutionCount: 1, + Data: map[string]interface{}{"text/plain": "earlier"}, + }), + kernelMessage(t, MsgExecuteReply, "earlier-reply", session, earlier, ExecuteReply{ExecutionCount: 1, Status: "ok"}), + kernelMessage(t, MsgStatus, "earlier-idle", session, earlier, StatusUpdate{ExecutionState: StateIdle}), + ) { + return + } + time.Sleep(50 * time.Millisecond) + if !writeMessages(t, conn, + kernelMessage(t, MsgStatus, "own-busy", session, own, StatusUpdate{ExecutionState: StateBusy}), + kernelMessage(t, MsgStream, "own-stream", session, own, StreamOutput{Name: StreamStdout, Text: "own\n"}), + kernelMessage(t, MsgExecuteReply, "own-reply", session, own, ExecuteReply{ExecutionCount: 2, Status: "ok"}), + kernelMessage(t, MsgStatus, "own-idle", session, own, StatusUpdate{ExecutionState: StateIdle}), + ) { + return + } + <-release + }) + defer server.Close() + defer close(release) + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/kernels/test-kernel-id/channels" + client := NewClient("", nil) + require.NoError(t, client.Connect(wsURL)) + defer client.Disconnect() + + resultChan := make(chan *ExecutionResult, 10) + require.NoError(t, client.ExecuteCodeStream("print('own')", resultChan)) + + var streams []string + completed := false + for _, result := range collectResults(t, resultChan) { + require.Nil(t, result.Error, "an earlier request's error must not be reported for this run") + require.Nil(t, result.ExecutionData, "an earlier request's result must not be reported for this run") + for _, stream := range result.Stream { + streams = append(streams, stream.Text) + } + if result.ExecutionTime > 0 { + completed = true + } + } + require.Equal(t, []string{"own\n"}, streams) + require.True(t, completed, "expected the run to complete on its own idle and reply") +} + +// TestExecuteCodeStreamIgnoresReplyOfAnotherRequest covers an earlier request's execute_reply, +// which can arrive at any point of the run. Taken as the run's own, its execution count would end +// finalization before the run's own result, and its error would make a later kernel stop look +// like a run that had already replied, so the KernelStopped error would never be sent. +func TestExecuteCodeStreamIgnoresReplyOfAnotherRequest(t *testing.T) { + previousPollInterval := execdflag.JupyterIdlePollInterval + execdflag.JupyterIdlePollInterval = time.Millisecond + t.Cleanup(func() { + execdflag.JupyterIdlePollInterval = previousPollInterval + }) + + earlier := Header{MessageID: "earlier-execute-request-id", MessageType: string(MsgExecuteRequest)} + cases := []struct { + name string + // messages are sent after the run's own busy status. + messages func(t *testing.T, session string, own Header) [][]Message + check func(t *testing.T, results []*ExecutionResult) + }{ + { + name: "earlier reply while finalizing", + messages: func(t *testing.T, session string, own Header) [][]Message { + return [][]Message{ + { + kernelMessage(t, MsgStatus, "own-idle", session, own, StatusUpdate{ExecutionState: StateIdle}), + kernelMessage(t, MsgExecuteReply, "earlier-reply", session, earlier, ExecuteReply{ExecutionCount: 1, Status: "ok"}), + }, + { + kernelMessage(t, MsgExecuteResult, "own-result", session, own, ExecuteResult{ + ExecutionCount: 2, + Data: map[string]interface{}{"text/plain": "own"}, + }), + kernelMessage(t, MsgExecuteReply, "own-reply", session, own, ExecuteReply{ExecutionCount: 2, Status: "ok"}), + }, + } + }, + check: func(t *testing.T, results []*ExecutionResult) { + var own bool + for _, result := range results { + if result.ExecutionData != nil { + require.Equal(t, "own", result.ExecutionData["text/plain"]) + own = true + } + } + require.True(t, own, "the run's own result must arrive before the stream closes") + }, + }, + { + name: "earlier error reply, then the kernel dies", + messages: func(t *testing.T, session string, _ Header) [][]Message { + dead, _ := json.Marshal(StatusUpdate{ExecutionState: StateDead}) + return [][]Message{ + { + kernelMessage(t, MsgExecuteReply, "earlier-reply", session, earlier, ExecuteReply{ + ExecutionCount: 1, + Status: "error", + ErrorOutput: ErrorOutput{EName: "EarlierError", EValue: "earlier"}, + }), + }, + { + {Header: Header{MessageID: "dead-id", Session: session, MessageType: string(MsgStatus)}, Content: dead, Channel: "iopub"}, + }, + } + }, + check: func(t *testing.T, results []*ExecutionResult) { + var stopped *ErrorOutput + for _, result := range results { + if result.Error != nil { + stopped = result.Error + } + } + require.NotNil(t, stopped, "expected the KernelStopped error before the stream closed") + require.Equal(t, "KernelStopped", stopped.EName) + }, + }, + } + + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + release := make(chan struct{}) + server := createTestServer(t, func(conn *websocket.Conn) { + var executeRequest Message + if err := conn.ReadJSON(&executeRequest); err != nil { + t.Errorf("failed to read execution request: %v", err) + return + } + session := executeRequest.Header.Session + own := executeRequest.Header + if !writeMessages(t, conn, kernelMessage(t, MsgStatus, "own-busy", session, own, StatusUpdate{ExecutionState: StateBusy})) { + return + } + for _, batch := range tc.messages(t, session, own) { + if !writeMessages(t, conn, batch...) { + return + } + // Many poll intervals between batches. + time.Sleep(20 * time.Millisecond) + } + <-release + }) + defer server.Close() + defer close(release) + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/kernels/test-kernel-id/channels" + client := NewClient("", nil) + require.NoError(t, client.Connect(wsURL)) + defer client.Disconnect() + + resultChan := make(chan *ExecutionResult, 10) + require.NoError(t, client.ExecuteCodeStream("x", resultChan)) + tc.check(t, collectResults(t, resultChan)) + }) + } +} + +func setKernelWake(t *testing.T, delay time.Duration, attempts int) { + t.Helper() + previousDelay, previousAttempts := kernelWakeDelay, kernelWakeAttempts + kernelWakeDelay, kernelWakeAttempts = delay, attempts + t.Cleanup(func() { + kernelWakeDelay, kernelWakeAttempts = previousDelay, previousAttempts + }) +} + +// TestExecuteCodeStreamWakesKernelThatHasNotPickedUpTheRequest covers a kernel that leaves the +// run's request unread on its shell socket until something new connects to the kernel. The +// server here holds the request until a second websocket arrives, as such a kernel does; the run +// must open one and then end normally. +func TestExecuteCodeStreamWakesKernelThatHasNotPickedUpTheRequest(t *testing.T) { + setKernelWake(t, 20*time.Millisecond, 4) + + var connections atomic.Int32 + woken := make(chan struct{}) + var wakeOnce sync.Once + release := make(chan struct{}) + server := createTestServer(t, func(conn *websocket.Conn) { + if connections.Add(1) > 1 { + wakeOnce.Do(func() { close(woken) }) + return + } + var executeRequest Message + if err := conn.ReadJSON(&executeRequest); err != nil { + t.Errorf("failed to read execution request: %v", err) + return + } + select { + case <-woken: + case <-release: + return + } + session := executeRequest.Header.Session + own := executeRequest.Header + if !writeMessages(t, conn, + kernelMessage(t, MsgStatus, "own-busy", session, own, StatusUpdate{ExecutionState: StateBusy}), + kernelMessage(t, MsgExecuteReply, "own-reply", session, own, ExecuteReply{ExecutionCount: 1, Status: "ok"}), + kernelMessage(t, MsgStatus, "own-idle", session, own, StatusUpdate{ExecutionState: StateIdle}), + ) { + return + } + <-release + }) + defer server.Close() + defer close(release) + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/kernels/test-kernel-id/channels" + client := NewClient("", nil) + require.NoError(t, client.Connect(wsURL)) + defer client.Disconnect() + + resultChan := make(chan *ExecutionResult, 10) + require.NoError(t, client.ExecuteCodeStream("x = 1", resultChan)) + + completed := false + for _, result := range collectResults(t, resultChan) { + if result.ExecutionTime > 0 { + completed = true + } + } + require.True(t, completed) + require.GreaterOrEqual(t, connections.Load(), int32(2), "expected a second connection to wake the kernel") +} + +// TestExecuteCodeStreamDoesNotWakeKernelThatPickedUpTheRequest: once any message for the request +// arrives the kernel is working on it, however long it takes, and no extra connection is made. +func TestExecuteCodeStreamDoesNotWakeKernelThatPickedUpTheRequest(t *testing.T) { + setKernelWake(t, 10*time.Millisecond, 4) + + var connections atomic.Int32 + release := make(chan struct{}) + server := createTestServer(t, func(conn *websocket.Conn) { + if connections.Add(1) > 1 { + return + } + var executeRequest Message + if err := conn.ReadJSON(&executeRequest); err != nil { + t.Errorf("failed to read execution request: %v", err) + return + } + session := executeRequest.Header.Session + own := executeRequest.Header + if !writeMessages(t, conn, kernelMessage(t, MsgStatus, "own-busy", session, own, StatusUpdate{ExecutionState: StateBusy})) { + return + } + // Longer than every wake-up attempt put together. + time.Sleep(200 * time.Millisecond) + if !writeMessages(t, conn, + kernelMessage(t, MsgExecuteReply, "own-reply", session, own, ExecuteReply{ExecutionCount: 1, Status: "ok"}), + kernelMessage(t, MsgStatus, "own-idle", session, own, StatusUpdate{ExecutionState: StateIdle}), + ) { + return + } + <-release + }) + defer server.Close() + defer close(release) + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/kernels/test-kernel-id/channels" + client := NewClient("", nil) + require.NoError(t, client.Connect(wsURL)) + defer client.Disconnect() + + resultChan := make(chan *ExecutionResult, 10) + require.NoError(t, client.ExecuteCodeStream("import time; time.sleep(0.2)", resultChan)) + collectResults(t, resultChan) + require.Equal(t, int32(1), connections.Load(), "a kernel that picked up the request must not be woken") +} + +// TestExecuteCodeStreamStopsWakingKernel covers the two ways a run can finish with nothing for its +// request ever arriving: its kernel stopped, which ends the run, or the run gave up and closed its +// connection (runJupyterCode does so when its context is cancelled). Waking must stop with it: a +// wake-up after that would reach the same kernel during the next run, and the kernel's reply to +// the extra kernel_info_request can itself meet the next run's request. +func TestExecuteCodeStreamStopsWakingKernel(t *testing.T) { + cases := []struct { + name string + // kernelStops makes the server end the run with a dead status; otherwise the client + // disconnects. + kernelStops bool + }{ + {name: "run ended", kernelStops: true}, + {name: "connection closed", kernelStops: false}, + } + for _, tc := range cases { + t.Run(tc.name, func(t *testing.T) { + setKernelWake(t, 50*time.Millisecond, 3) + + var connections atomic.Int32 + requested := make(chan struct{}) + release := make(chan struct{}) + server := createTestServer(t, func(conn *websocket.Conn) { + if connections.Add(1) > 1 { + return + } + var executeRequest Message + if err := conn.ReadJSON(&executeRequest); err != nil { + t.Errorf("failed to read execution request: %v", err) + return + } + close(requested) + if tc.kernelStops { + dead, _ := json.Marshal(StatusUpdate{ExecutionState: StateDead}) + if !writeMessages(t, conn, Message{ + Header: Header{MessageID: "dead-id", Session: executeRequest.Header.Session, MessageType: string(MsgStatus)}, + Content: dead, + }) { + return + } + } + <-release + }) + defer server.Close() + defer close(release) + + wsURL := "ws" + strings.TrimPrefix(server.URL, "http") + "/api/kernels/test-kernel-id/channels" + client := NewClient("", nil) + require.NoError(t, client.Connect(wsURL)) + defer client.Disconnect() + + resultChan := make(chan *ExecutionResult, 10) + require.NoError(t, client.ExecuteCodeStream("x", resultChan)) + <-requested + if tc.kernelStops { + collectResults(t, resultChan) + } else { + client.Disconnect() + } + // Longer than every wake-up attempt put together (50 + 100 + 200 ms). + time.Sleep(500 * time.Millisecond) + require.Equal(t, int32(1), connections.Load(), "no wake-up may follow the end of the run or its connection") + }) + } +} + +// TestReceiveMessagesReadsOnlyItsOwnConnection covers a reader that is still handling a message +// when its run disconnects and the next run connects. It must stop with its own connection, not +// go on to read the new one alongside that connection's own reader. +func TestReceiveMessagesReadsOnlyItsOwnConnection(t *testing.T) { + serversDone := make(chan struct{}) + first := createTestServer(t, func(conn *websocket.Conn) { + if !writeMessages(t, conn, kernelMessage(t, MsgStream, "first-stream", "session", Header{}, StreamOutput{Name: StreamStdout, Text: "first\n"})) { + return + } + <-serversDone + }) + defer first.Close() + second := createTestServer(t, func(_ *websocket.Conn) { + <-serversDone + }) + defer second.Close() + defer close(serversDone) + + client := NewClient("", nil) + handling := make(chan struct{}) + unblock := make(chan struct{}) + var handlingOnce sync.Once + client.registerHandler(MsgStream, func(*Message) { + handlingOnce.Do(func() { close(handling) }) + <-unblock + }) + + // Start the first reader here rather than through Connect, so the test can see it return. + firstURL := "ws" + strings.TrimPrefix(first.URL, "http") + "/api/kernels/test-kernel-id/channels" + conn, resp, err := websocket.DefaultDialer.Dial(firstURL, nil) + if resp != nil && resp.Body != nil { + resp.Body.Close() + } + require.NoError(t, err) + client.mu.Lock() + client.conn = conn + client.connClosed = make(chan struct{}) + client.mu.Unlock() + exited := make(chan struct{}) + go func() { + client.receiveMessages(conn) + close(exited) + }() + + select { + case <-handling: + case <-time.After(2 * time.Second): + t.Fatal("the first reader never handled the first server's message") + } + client.Disconnect() + secondURL := "ws" + strings.TrimPrefix(second.URL, "http") + "/api/kernels/test-kernel-id/channels" + require.NoError(t, client.Connect(secondURL)) + defer client.Disconnect() + close(unblock) + + select { + case <-exited: + case <-time.After(time.Second): + t.Fatal("the first reader went on reading after its connection was closed") + } +}