on main help sign in
fix(server): probe readiness over HTTP instead of statting the socket
9b366c3Christian Rocha
internal/client/client.go+8
@@ -121,6 +121,14 @@ func (c *Client) ShutdownServer(ctx context.Context) error {
121 return nil
121 return nil
122}
122}
123
123
124+// Dial opens a connection to the server using the same scheme-aware
125+// logic the client uses for its HTTP transport. Exposed so callers can
126+// reuse the dialer when they need to construct sibling HTTP transports
127+// (e.g. a readiness probe in the CLI).
128+func (c *Client) Dial(ctx context.Context, network, address string) (net.Conn, error) {
129+ return c.dialer(ctx, network, address)
130+}
131+
124func (c *Client) dialer(ctx context.Context, network, address string) (net.Conn, error) {
132func (c *Client) dialer(ctx context.Context, network, address string) (net.Conn, error) {
125 d := net.Dialer{
133 d := net.Dialer{
126 Timeout: 30 * time.Second,
134 Timeout: 30 * time.Second,
internal/cmd/root.go+112−28
@@ -9,6 +9,8 @@ import (
9 "io"
9 "io"
10 "io/fs"
10 "io/fs"
11 "log/slog"
11 "log/slog"
12+ "net"
13+ "net/http"
12 "net/url"
14 "net/url"
13 "os"
15 "os"
14 "os/exec"
16 "os/exec"
@@ -356,22 +358,7 @@ func connectToServer(cmd *cobra.Command) (*client.Client, *proto.Workspace, func
356
358
357 ws, err := c.CreateWorkspace(ctx, wsReq)
359 ws, err := c.CreateWorkspace(ctx, wsReq)
358 if err != nil {
360 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)
361+ return nil, nil, nil, fmt.Errorf("failed to create workspace: %v", err)
374− }
375 }
362 }
376
363
377 if shouldEnableMetrics(ws.Config) {
364 if shouldEnableMetrics(ws.Config) {
@@ -410,23 +397,120 @@ func ensureServer(cmd *cobra.Command, hostURL *url.URL) error {
410 }
397 }
411 }
398 }
412
399
413− var err error
400+ if err := waitForServerReady(cmd.Context(), hostURL); err != nil {
414− for range 10 {
401+ return fmt.Errorf("failed to initialize crush server: %v", err)
415− _, err = os.Stat(hostURL.Host)
402+ }
416− if err == nil {
403+ }
417− break
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
418− }
415+ }
419− select {
416+ d, err := time.ParseDuration(v)
420− case <-cmd.Context().Done():
417+ if err != nil || d <= 0 {
421− return cmd.Context().Err()
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+
422− case <-time.After(100 * time.Millisecond):
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
423 }
446 }
447+ return fmt.Errorf("timed out waiting for server readiness")
424 }
448 }
449+
450+ attemptCtx, cancel := context.WithTimeout(ctx, perAttempt)
451+ err := probeHealth(attemptCtx, httpClient, reqURL, hostURL)
452+ cancel()
425− if err != nil {
453+ if err == nil {
426− return fmt.Errorf("failed to initialize crush server: %v", err)
454+ return nil
427 }
455 }
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
428 }
481 }
429
482
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+ }
430 return nil
514 return nil
431}
515}
432
516