feat: liveness and readiness probes (#6)

Signed-off-by: Carlos Alexandro Becker <[email protected]>

internal/proxy/proxy.go

@@ -30,9 +30,13 @@ type Proxy struct {
 	mu      sync.Mutex
 	servers []*Server
 
+	initialized  atomic.Bool
 	shuttingDown atomic.Bool
 }
 
+func (p *Proxy) Ready() bool { return p.initialized.Load() }
+func (p *Proxy) Live() bool  { return !p.shuttingDown.Load() && p.initialized.Load() }
+
 func (p *Proxy) Stats() []ServerStat {
 	var result []ServerStat
 	for _, s := range p.servers {
@@ -120,6 +124,7 @@ func (p *Proxy) update(rpcs []seed.RPC) error {
 	})
 
 	slog.Info("updated server list", "total", len(p.servers))
+	p.initialized.Store(true)
 	return nil
 }
 

main.go

@@ -43,6 +43,16 @@ func main() {
 	indexTpl := template.Must(template.New("stats").Parse(string(index)))
 
 	m := http.NewServeMux()
+	m.Handle("/health/ready", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !proxyHandler.Ready() {
+			w.WriteHeader(http.StatusServiceUnavailable)
+		}
+	}))
+	m.Handle("/health/live", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		if !proxyHandler.Live() {
+			w.WriteHeader(http.StatusServiceUnavailable)
+		}
+	}))
 	m.Handle("/rpc", proxyHandler)
 	m.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 		if err := indexTpl.Execute(w, proxyHandler.Stats()); err != nil {