diff --git a/runner/internal/runner/api/http.go b/runner/internal/runner/api/http.go index 4919852aff..d77690f645 100644 --- a/runner/internal/runner/api/http.go +++ b/runner/internal/runner/api/http.go @@ -188,7 +188,7 @@ func (s *Server) pullGetHandler(w http.ResponseWriter, r *http.Request) (interfa } if s.executor.GetRunnerState() == executor.WaitLogsFinished { - defer func() { close(s.pullDoneCh) }() + defer s.closePullDone() } return s.executor.GetHistory(timestamp), nil } diff --git a/runner/internal/runner/api/server.go b/runner/internal/runner/api/server.go index 227ea1dbb3..72de48c138 100644 --- a/runner/internal/runner/api/server.go +++ b/runner/internal/runner/api/server.go @@ -5,6 +5,7 @@ import ( "errors" "net/http" _ "net/http/pprof" + "sync" "time" "github.com/dstackai/dstack/runner/internal/common/api" @@ -19,7 +20,9 @@ type Server struct { shutdownCh chan interface{} // server closes this chan on shutdown jobBarrierCh chan interface{} // only server listens on this chan pullDoneCh chan interface{} // Closed then /api/pull gave everything + pullDoneOnce sync.Once wsDoneCh chan interface{} // Closed then /logs_ws gave everything + wsDoneOnce sync.Once startWaitDuration time.Duration logsWaitDuration time.Duration @@ -123,6 +126,23 @@ loop: return nil } +// closePullDone reports that /api/pull has served the final logs. +// +// More than one request may observe the WaitLogsFinished state: the state is set as soon as +// the job is asked to stop, while the server keeps serving until the executor returns, which +// may take arbitrarily long if the job leaves processes behind. Closing must be idempotent. +func (s *Server) closePullDone() { + s.pullDoneOnce.Do(func() { close(s.pullDoneCh) }) +} + +// closeWsDone reports that a /logs_ws stream has sent the final logs. +// +// Nothing limits the number of concurrent connections, and each one is served by its own +// goroutine, so more than one may drain. Closing must be idempotent. +func (s *Server) closeWsDone() { + s.wsDoneOnce.Do(func() { close(s.wsDoneCh) }) +} + func (s *Server) stop() { s.executor.Lock() defer s.executor.Unlock() diff --git a/runner/internal/runner/api/server_test.go b/runner/internal/runner/api/server_test.go new file mode 100644 index 0000000000..1bbb5e34df --- /dev/null +++ b/runner/internal/runner/api/server_test.go @@ -0,0 +1,165 @@ +package api + +import ( + "context" + "io" + "net/http" + "net/http/httptest" + "strings" + "sync" + "testing" + "time" + + "github.com/gorilla/websocket" + "github.com/stretchr/testify/require" + + "github.com/dstackai/dstack/runner/internal/common/types" + "github.com/dstackai/dstack/runner/internal/runner/executor" + "github.com/dstackai/dstack/runner/internal/runner/schemas" +) + +// fakeExecutor implements executor.Executor with the minimum needed to drive the handlers. +type fakeExecutor struct { + mu sync.RWMutex + state string +} + +func (e *fakeExecutor) SetJob(schemas.SubmitBody) {} +func (e *fakeExecutor) WriteFileArchive(string, io.Reader) error { return nil } +func (e *fakeExecutor) WriteRepoBlob(io.Reader) error { return nil } +func (e *fakeExecutor) Run(context.Context) error { return nil } + +func (e *fakeExecutor) GetHistory(int64) *schemas.PullResponse { return &schemas.PullResponse{} } +func (e *fakeExecutor) GetJobWsLogsHistory() []schemas.LogEvent { return nil } + +func (e *fakeExecutor) GetRunnerState() string { return e.state } +func (e *fakeExecutor) SetRunnerState(state string) { e.state = state } + +func (e *fakeExecutor) GetJobInfo(context.Context) (string, string, error) { return "", "", nil } +func (e *fakeExecutor) SetJobState(context.Context, schemas.JobState) {} +func (e *fakeExecutor) SetJobStateWithTerminationReason( + context.Context, schemas.JobState, types.TerminationReason, string, +) { +} + +func (e *fakeExecutor) Lock() { e.mu.Lock() } +func (e *fakeExecutor) Unlock() { e.mu.Unlock() } +func (e *fakeExecutor) RLock() { e.mu.RLock() } +func (e *fakeExecutor) RUnlock() { e.mu.RUnlock() } + +func newTestServer(t *testing.T, state string) *Server { + t.Helper() + s, err := NewServer(t.Context(), "localhost:0", "test", &fakeExecutor{state: state}) + require.NoError(t, err) + return s +} + +func pull(t *testing.T, s *Server) { + t.Helper() + _, err := s.pullGetHandler( + httptest.NewRecorder(), + httptest.NewRequest(http.MethodGet, "/api/pull", nil), + ) + require.NoError(t, err) +} + +// The runner enters WaitLogsFinished as soon as the job is asked to stop, but keeps serving +// until the executor returns, which never happens while the job holds the pty open. The dstack +// server goes on polling, so several pulls observe the state. +func TestPullGetHandler_RepeatedFinalPulls(t *testing.T) { + s := newTestServer(t, executor.WaitLogsFinished) + + for range 3 { + pull(t, s) + } + + select { + case <-s.pullDoneCh: + default: + t.Fatal("pullDoneCh must be closed after a pull in the WaitLogsFinished state") + } +} + +// The handler takes only a read lock, so pulls can observe the state concurrently. +func TestPullGetHandler_ConcurrentFinalPulls(t *testing.T) { + s := newTestServer(t, executor.WaitLogsFinished) + + var wg sync.WaitGroup + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + pull(t, s) + }() + } + wg.Wait() + + select { + case <-s.pullDoneCh: + default: + t.Fatal("pullDoneCh must be closed after a pull in the WaitLogsFinished state") + } +} + +func TestPullGetHandler_NotFinishedKeepsPullDoneOpen(t *testing.T) { + s := newTestServer(t, executor.ServeLogs) + + pull(t, s) + + select { + case <-s.pullDoneCh: + t.Fatal("pullDoneCh must stay open while the executor still serves logs") + default: + } +} + +// Nothing limits the number of concurrent /logs_ws connections, and each is drained by its +// own goroutine. +func TestCloseWsDone_Idempotent(t *testing.T) { + s := newTestServer(t, executor.WaitLogsFinished) + + var wg sync.WaitGroup + for range 8 { + wg.Add(1) + go func() { + defer wg.Done() + s.closeWsDone() + }() + } + wg.Wait() + + select { + case <-s.wsDoneCh: + default: + t.Fatal("wsDoneCh must be closed") + } +} + +// Two attached clients -- `dstack apply` in one terminal and `dstack attach` in another -- +// each open their own /logs_ws stream. Both drain, both see shutdownCh, and both reach the +// close. Unlike the /api/pull handler, streamJobLogs runs in a bare goroutine, so an +// unrecovered panic there takes down the whole runner. +func TestLogsWs_TwoClientsBothDrain(t *testing.T) { + s := newTestServer(t, executor.ServeLogs) + httpSrv := httptest.NewServer(s.srv.Handler) + defer httpSrv.Close() + + wsURL := "ws" + strings.TrimPrefix(httpSrv.URL, "http") + "/logs_ws" + for range 2 { + conn, _, err := websocket.DefaultDialer.Dial(wsURL, nil) + require.NoError(t, err) + defer func() { _ = conn.Close() }() + } + + // Let both streams reach the drained-and-sleeping branch of their loop. + time.Sleep(300 * time.Millisecond) + close(s.shutdownCh) + + select { + case <-s.wsDoneCh: + case <-time.After(5 * time.Second): + t.Fatal("wsDoneCh must be closed once a stream has drained after shutdown") + } + // Give the second stream time to reach its own close. + time.Sleep(500 * time.Millisecond) +} diff --git a/runner/internal/runner/api/ws.go b/runner/internal/runner/api/ws.go index 3229701a68..0df24f2231 100644 --- a/runner/internal/runner/api/ws.go +++ b/runner/internal/runner/api/ws.go @@ -81,7 +81,7 @@ func (s *Server) streamJobLogs(ctx context.Context, conn *websocket.Conn, params case <-s.shutdownCh: if currentPos >= len(jobLogsWsHistory) { s.executor.RUnlock() - close(s.wsDoneCh) + s.closeWsDone() return } default: