Proxy health check (#381)

superfly · Aug 11, 2023 · b135766 · b135766
1 parent 91f35ae
commit b135766
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 5 deletions.
diff --git a/cmd/litefs/config.go b/cmd/litefs/config.go
@@ -57,6 +57,8 @@ func NewConfig() Config {
 
 	config.Log.Format = "text"
 
+	config.Proxy.MaxLag = http.DefaultMaxLag
+
 	config.Tracing.Enabled = true
 	config.Tracing.MaxSize = DefaultTracingMaxSize
 	config.Tracing.MaxCount = DefaultTracingMaxCount
@@ -110,11 +112,12 @@ type HTTPConfig struct {
 
 // ProxyConfig represents the configuration for the HTTP proxy server.
 type ProxyConfig struct {
-	Addr        string   `yaml:"addr"`
-	Target      string   `yaml:"target"`
-	DB          string   `yaml:"db"`
-	Debug       bool     `yaml:"debug"`
-	Passthrough []string `yaml:"passthrough"`
+	Addr        string        `yaml:"addr"`
+	Target      string        `yaml:"target"`
+	DB          string        `yaml:"db"`
+	MaxLag      time.Duration `yaml:"max-lag"`
+	Debug       bool          `yaml:"debug"`
+	Passthrough []string      `yaml:"passthrough"`
 }
 
 // LeaseConfig represents a generic configuration for all lease types.

diff --git a/cmd/litefs/mount_linux.go b/cmd/litefs/mount_linux.go
@@ -503,6 +503,7 @@ func (c *MountCommand) runProxyServer(ctx context.Context) error {
 	server.Target = c.Config.Proxy.Target
 	server.DBName = c.Config.Proxy.DB
 	server.Addr = c.Config.Proxy.Addr
+	server.MaxLag = c.Config.Proxy.MaxLag
 	server.Debug = c.Config.Proxy.Debug
 	server.Passthroughs = passthroughs
 	if err := server.Listen(); err != nil {

diff --git a/http/proxy_server.go b/http/proxy_server.go
@@ -24,6 +24,8 @@ const (
 	DefaultPollTXIDInterval = 1 * time.Millisecond
 	DefaultPollTXIDTimeout  = 5 * time.Second
 
+	DefaultMaxLag = 10 * time.Second
+
 	DefaultCookieExpiry = 5 * time.Minute
 )
 
@@ -61,6 +63,9 @@ type ProxyServer struct {
 	PollTXIDInterval time.Duration
 	PollTXIDTimeout  time.Duration
 
+	// Maximum allowable lag before the health endpoint returns an error code.
+	MaxLag time.Duration
+
 	// Time before cookie expires on client.
 	CookieExpiry time.Duration
 
@@ -74,6 +79,7 @@ func NewProxyServer(store *litefs.Store) *ProxyServer {
 
 		PollTXIDInterval: DefaultPollTXIDInterval,
 		PollTXIDTimeout:  DefaultPollTXIDTimeout,
+		MaxLag:           DefaultMaxLag,
 		CookieExpiry:     DefaultCookieExpiry,
 	}
 
@@ -167,6 +173,12 @@ func (s *ProxyServer) serveHTTP(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	// Handle health check endpoint.
+	if r.Method == http.MethodGet && r.URL.Path == "/litefs/healthz" {
+		s.serveGetHealthz(w, r)
+		return
+	}
+
 	switch r.Method {
 	case http.MethodGet:
 		s.serveRead(w, r)
@@ -177,6 +189,18 @@ func (s *ProxyServer) serveHTTP(w http.ResponseWriter, r *http.Request) {
 	}
 }
 
+func (s *ProxyServer) serveGetHealthz(w http.ResponseWriter, r *http.Request) {
+	lag := s.store.Lag()
+	if s.MaxLag > 0 && lag > s.MaxLag {
+		s.logf("proxy: %s %s: current replication lag of %s exceeds maximum threshold of %s", r.Method, r.URL.Path, lag, s.MaxLag)
+		http.Error(w, "Replication lag exceeded", http.StatusServiceUnavailable)
+		return
+	}
+
+	w.WriteHeader(http.StatusOK)
+	_, _ = w.Write([]byte("OK\n"))
+}
+
 func (s *ProxyServer) serveRead(w http.ResponseWriter, r *http.Request) {
 	// Determine the last write TXID seen by
 	var txid ltx.TXID