@@ -9,6 +9,8 @@ import (
99 "io"
1010 "io/fs"
1111 "log/slog"
12+ "net"
13+ "net/http"
1214 "net/url"
1315 "os"
1416 "os/exec"
@@ -356,22 +358,7 @@ func connectToServer(cmd *cobra.Command) (*client.Client, *proto.Workspace, func
356358
357359 ws, err := c.CreateWorkspace(ctx, wsReq)
358360 if err != nil {
359− // The server socket may exist before the HTTP handler is ready.
360− // Retry a few times with a short backoff.
361− for range 5 {
362− select {
363− case <-ctx.Done():
364− return nil, nil, nil, ctx.Err()
365− case <-time.After(200 * time.Millisecond):
366− }
367− ws, err = c.CreateWorkspace(ctx, wsReq)
368− if err == nil {
369− break
370− }
371− }
372− if err != nil {
373− return nil, nil, nil, fmt.Errorf("failed to create workspace: %v", err)
374− }
361+ return nil, nil, nil, fmt.Errorf("failed to create workspace: %v", err)
375362 }
376363
377364 if shouldEnableMetrics(ws.Config) {
@@ -410,23 +397,120 @@ func ensureServer(cmd *cobra.Command, hostURL *url.URL) error {
410397 }
411398 }
412399
413− var err error
414− for range 10 {
415− _, err = os.Stat(hostURL.Host)
416− if err == nil {
417− break
418− }
419− select {
420− case <-cmd.Context().Done():
421− return cmd.Context().Err()
422− case <-time.After(100 * time.Millisecond):
400+ if err := waitForServerReady(cmd.Context(), hostURL); err != nil {
401+ return fmt.Errorf("failed to initialize crush server: %v", err)
402+ }
403+ }
404+
405+ return nil
406+}
407+
408+// serverReadyTimeout returns the total budget for the readiness probe.
409+// Overridable via CRUSH_SERVER_READY_TIMEOUT (parsed as a Go duration).
410+func serverReadyTimeout() time.Duration {
411+ const def = 10 * time.Second
412+ v := os.Getenv("CRUSH_SERVER_READY_TIMEOUT")
413+ if v == "" {
414+ return def
415+ }
416+ d, err := time.ParseDuration(v)
417+ if err != nil || d <= 0 {
418+ return def
419+ }
420+ return d
421+}
422+
423+// waitForServerReady polls GET /v1/health until the server responds with
424+// any 2xx status or the total timeout elapses. Each attempt uses a short
425+// per-attempt timeout so a hung listener doesn't burn the whole budget.
426+//
427+// The HTTP transport is built to mirror how *client.Client dials so the
428+// same unix socket / npipe / tcp setups all work uniformly here.
429+func waitForServerReady(ctx context.Context, hostURL *url.URL) error {
430+ httpClient, reqURL, err := readinessHTTPClient(hostURL)
431+ if err != nil {
432+ return err
433+ }
434+
435+ const perAttempt = 100 * time.Millisecond
436+ deadline := time.Now().Add(serverReadyTimeout())
437+
438+ var lastErr error
439+ for {
440+ if err := ctx.Err(); err != nil {
441+ return err
442+ }
443+ if time.Now().After(deadline) {
444+ if lastErr != nil {
445+ return lastErr
423446 }
447+ return fmt.Errorf("timed out waiting for server readiness")
424448 }
425− if err != nil {
426− return fmt.Errorf("failed to initialize crush server: %v", err)
449+
450+ attemptCtx, cancel := context.WithTimeout(ctx, perAttempt)
451+ err := probeHealth(attemptCtx, httpClient, reqURL, hostURL)
452+ cancel()
453+ if err == nil {
454+ return nil
427455 }
456+ lastErr = err
457+
458+ select {
459+ case <-ctx.Done():
460+ return ctx.Err()
461+ case <-time.After(perAttempt):
462+ }
463+ }
464+}
465+
466+// readinessHTTPClient builds an *http.Client whose transport dials the
467+// server using the same scheme-aware logic as *client.Client (unix
468+// socket, named pipe, or tcp).
469+func readinessHTTPClient(hostURL *url.URL) (*http.Client, string, error) {
470+ c, err := client.NewClient("", hostURL.Scheme, hostURL.Host)
471+ if err != nil {
472+ return nil, "", err
473+ }
474+
475+ tr := http.DefaultTransport.(*http.Transport).Clone()
476+ tr.DialContext = func(ctx context.Context, network, addr string) (net.Conn, error) {
477+ return c.Dial(ctx, network, addr)
478+ }
479+ if hostURL.Scheme == "unix" || hostURL.Scheme == "npipe" {
480+ tr.DisableCompression = true
428481 }
429482
483+ httpClient := &http.Client{Transport: tr}
484+
485+ // For unix sockets / named pipes we still need a syntactically valid
486+ // HTTP URL; the actual address is resolved by the dialer.
487+ host := hostURL.Host
488+ if hostURL.Scheme == "unix" || hostURL.Scheme == "npipe" {
489+ host = client.DummyHost
490+ }
491+ reqURL := (&url.URL{Scheme: "http", Host: host, Path: "/v1/health"}).String()
492+ return httpClient, reqURL, nil
493+}
494+
495+// probeHealth issues a single GET to the readiness endpoint and treats
496+// any 2xx response as success.
497+func probeHealth(ctx context.Context, h *http.Client, reqURL string, hostURL *url.URL) error {
498+ req, err := http.NewRequestWithContext(ctx, http.MethodGet, reqURL, nil)
499+ if err != nil {
500+ return err
501+ }
502+ if hostURL.Scheme == "unix" || hostURL.Scheme == "npipe" {
503+ req.Host = client.DummyHost
504+ }
505+ rsp, err := h.Do(req)
506+ if err != nil {
507+ return err
508+ }
509+ defer rsp.Body.Close()
510+ _, _ = io.Copy(io.Discard, rsp.Body)
511+ if rsp.StatusCode < 200 || rsp.StatusCode >= 300 {
512+ return fmt.Errorf("server health check failed: %s", rsp.Status)
513+ }
430514 return nil
431515}
432516