feat: add health check

Author: mchristopher

main.go

@@ -70,6 +70,11 @@ func NewSegmentReverseProxy(cdn *url.URL, trackingAPI *url.URL, urlPrefix string
 	return &httputil.ReverseProxy{Director: director}
 }
 
+func healthHandler(w http.ResponseWriter, r *http.Request) {
+	w.WriteHeader(http.StatusOK)
+	w.Write([]byte("ready"))
+}
+
 var port = flag.String("port", "8080", "bind address")
 var debug = flag.Bool("debug", false, "debug mode")
 var urlPrefix = flag.String("url-prefix", "/", "URL prefix to strip from requests (can also use URL_PREFIX env var)")
@@ -89,11 +94,16 @@ func main() {
 	if urlPrefixEnv != "" && *urlPrefix == "/" {
 		*urlPrefix = urlPrefixEnv
 	}
+
+	mux := http.NewServeMux()
+	mux.HandleFunc("/_health", healthHandler)
+
 	proxy := NewSegmentReverseProxy(cdnURL, trackingAPIURL, *urlPrefix)
 	if *debug {
 		proxy = handlers.LoggingHandler(os.Stdout, proxy)
 		log.Printf("serving proxy at port %v with prefix %s\n", *port, *urlPrefix)
 	}
 
-	log.Fatal(http.ListenAndServe(":"+*port, proxy))
+	mux.Handle("/", proxy)
+	log.Fatal(http.ListenAndServe(":"+*port, mux))
 }

main_test.go

@@ -81,6 +81,27 @@ func TestSegmentReverseProxy(t *testing.T) {
 	}
 }
 
+func TestHealthEndpoint(t *testing.T) {
+	mux := http.NewServeMux()
+	mux.HandleFunc("/_health", healthHandler)
+	server := httptest.NewServer(mux)
+	defer server.Close()
+	resp, err := http.Get(server.URL + "/_health")
+	if err != nil {
+		t.Fatalf("Failed to GET /_health: %v", err)
+	}
+	defer resp.Body.Close()
+	if resp.StatusCode != http.StatusOK {
+		t.Errorf("Expected status 200, got %d", resp.StatusCode)
+	}
+	buf := make([]byte, 16)
+	n, _ := resp.Body.Read(buf)
+	body := string(buf[:n])
+	if body != "ready" {
+		t.Errorf("Expected body 'ready', got %q", body)
+	}
+}
+
 func mustParseUrl(raw string) *url.URL {
 	u, err := url.Parse(raw)
 	if err != nil {