golangLAKEHOUSE/internal/shared/server.go

// Package shared provides common HTTP server bootstrap for every
// Lakehouse-Go service. Each cmd/<service> calls Run with its name,
// bind address, and a route-registration callback. The factory wires
// chi, slog, /health, and graceful shutdown identically across all
// five binaries — the place where uniformity beats per-service
// flexibility.
//
// G1+ note: when queryd needs to drain a cgo DuckDB handle on
// shutdown, the simple shared factory will need a per-service hook
// (an io.Closer slice or an OnShutdown callback). For G0 a plain
// chi.Router + http.Server.Shutdown(ctx) is sufficient.
package shared

import (
	"context"
	"encoding/json"
	"errors"
	"log/slog"
	"net"
	"net/http"
	"os"
	"os/signal"
	"syscall"
	"time"

	"github.com/go-chi/chi/v5"
	"github.com/go-chi/chi/v5/middleware"
)

// HealthResponse is the JSON shape returned by /health on every
// service. Service-specific status hooks can extend it post-G0.
type HealthResponse struct {
	Status  string `json:"status"`
	Service string `json:"service"`
}

// RegisterRoutes is the per-service callback that wires its own
// routes onto the shared router AFTER /health has been mounted.
type RegisterRoutes func(r chi.Router)

// Run boots a chi router with slog logging, the /health endpoint,
// and graceful-shutdown handling. Blocks until SIGINT/SIGTERM or a
// fatal listener error.
//
// The logger is constructed locally and used as the request-logging
// sink. Run does NOT mutate the global slog default — callers that
// want their own slog.Default() should set it before calling Run.
// (Per Kimi review #4: shared library functions shouldn't silently
// mutate package globals.)
//
// Refuses to bind a non-loopback address unless the
// LH_<SERVICE>_ALLOW_NONLOOPBACK=1 env is set — closes the accidental
// 0.0.0.0 deploy path for R-001 (queryd /sql is RCE-equivalent off
// loopback, but the gate applies to every binary uniformly).
func Run(serviceName, addr string, register RegisterRoutes) error {
	if err := requireLoopbackOrOverride(serviceName, addr); err != nil {
		return err
	}

	logger := slog.New(slog.NewJSONHandler(os.Stdout, &slog.HandlerOptions{
		Level: slog.LevelInfo,
	}))

	r := chi.NewRouter()
	r.Use(middleware.RequestID)
	r.Use(middleware.RealIP)
	r.Use(middleware.Recoverer)
	r.Use(slogRequest(logger))

	r.Get("/health", func(w http.ResponseWriter, _ *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		_ = json.NewEncoder(w).Encode(HealthResponse{
			Status:  "ok",
			Service: serviceName,
		})
	})

	if register != nil {
		register(r)
	}

	srv := &http.Server{
		Addr:              addr,
		Handler:           r,
		ReadHeaderTimeout: 10 * time.Second,
	}

	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
	defer stop()

	// Race-safe startup: bind the listener synchronously BEFORE
	// returning so a fast bind error (e.g. port already in use) is
	// surfaced as Run's return value rather than racing the select.
	// Per Opus + Qwen BLOCK #1: the prior pattern could drop bind
	// errors when ctx.Done already fired or a fast failure happened
	// during select setup.
	ln, err := newListener(srv.Addr)
	if err != nil {
		return err
	}

	errCh := make(chan error, 1)
	go func() {
		logger.Info("listening", "service", serviceName, "addr", addr)
		if err := srv.Serve(ln); err != nil && !errors.Is(err, http.ErrServerClosed) {
			errCh <- err
		}
		close(errCh)
	}()

	select {
	case <-ctx.Done():
		logger.Info("shutdown signal received", "service", serviceName)
	case err := <-errCh:
		if err != nil {
			return err
		}
		// Server exited cleanly without a signal (unlikely but possible).
		return nil
	}

	shutdownCtx, cancel := context.WithTimeout(context.Background(), 15*time.Second)
	defer cancel()
	if err := srv.Shutdown(shutdownCtx); err != nil {
		return err
	}

	// Drain errCh so a late error from the listener goroutine
	// surfaces as the return value instead of leaking. After Shutdown
	// the channel will close on graceful exit; if a real error
	// landed first we return it.
	if err := <-errCh; err != nil {
		return err
	}
	return nil
}

// newListener binds the TCP listener up-front so bind errors are
// returned synchronously to Run's caller. Extracted into its own
// function for testability + to keep Run readable.
func newListener(addr string) (net.Listener, error) {
	return net.Listen("tcp", addr)
}

// slogRequest returns a chi middleware that logs each request via slog.
// Replaces chi's default text logger so all log output stays JSON.
func slogRequest(logger *slog.Logger) func(http.Handler) http.Handler {
	return func(next http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			start := time.Now()
			ww := middleware.NewWrapResponseWriter(w, r.ProtoMajor)
			defer func() {
				logger.Info("http",
					"method", r.Method,
					"path", r.URL.Path,
					"status", ww.Status(),
					"dur_ms", time.Since(start).Milliseconds(),
					"req_id", middleware.GetReqID(r.Context()),
				)
			}()
			next.ServeHTTP(ww, r)
		})
	}
}