add liveness and readiness probes

bouncer/bouncer.go

@@ -34,6 +34,7 @@ type DecisionCache interface {
 	Sync(ctx context.Context) error
 	Size() int
 	GetOriginCounts() map[string]int
+	IsReady() bool
 }
 
 //go:generate mockgen -destination=mocks/mock_captcha_service.go -package=mocks github.com/kdwils/envoy-proxy-bouncer/bouncer CaptchaService
@@ -130,6 +131,13 @@ func (b *Bouncer) Metrics(ctx context.Context) error {
 	return b.MetricsService.Run(ctx, b.config.Bouncer.MetricsInterval)
 }
 
+func (b *Bouncer) IsReady() bool {
+	if b.DecisionCache == nil {
+		return false
+	}
+	return b.DecisionCache.IsReady()
+}
+
 func (b *Bouncer) incRemediationMetric(name, remediationType string) {
 	if b.MetricsService == nil {
 		return

bouncer/components/decision_cache.go

@@ -19,6 +19,7 @@ type DecisionCache struct {
 	decisions      *cache.Cache[models.Decision]
 	mu             *sync.RWMutex
 	MetricsService *crowdsec.MetricsService
+	syncComplete   bool
 }
 
 func NewDecisionCache(apiKey, apiURL, tickerInterval string, MetricsService *crowdsec.MetricsService) (*DecisionCache, error) {
@@ -92,6 +93,12 @@ func (dc *DecisionCache) Size() int {
 	return dc.decisions.Size()
 }
 
+func (dc *DecisionCache) IsReady() bool {
+	dc.mu.RLock()
+	defer dc.mu.RUnlock()
+	return dc.syncComplete
+}
+
 func (dc *DecisionCache) GetOriginCounts() map[string]int {
 	originCounts := make(map[string]int)
 	if dc.decisions == nil {
@@ -168,8 +175,14 @@ func (dc *DecisionCache) Sync(ctx context.Context) error {
 				dc.MetricsService.Inc(key, "active_decisions", "ip", map[string]string{
 					"origin": origin,
 				})
+			}
 
+			dc.mu.Lock()
+			if !dc.syncComplete {
+				dc.syncComplete = true
+				logger.Info("initial decision sync complete")
 			}
+			dc.mu.Unlock()
 		}
 	}
 }

charts/envoy-proxy-bouncer/templates/deployment.yaml

@@ -45,6 +45,20 @@ spec:
               containerPort: {{ .Values.config.server.httpPort }}
               protocol: TCP
             {{- end }}
+          livenessProbe:
+            grpc:
+              port: {{ .Values.config.server.grpcPort }}
+            initialDelaySeconds: 10
+            periodSeconds: 10
+            timeoutSeconds: 5
+            failureThreshold: 3
+          readinessProbe:
+            grpc:
+              port: {{ .Values.config.server.grpcPort }}
+            initialDelaySeconds: 5
+            periodSeconds: 5
+            timeoutSeconds: 3
+            failureThreshold: 3
           env:
           {{- with .Values.config }}
             {{- if and .bouncer.apiKeySecretRef .bouncer.apiKeySecretRef.name .bouncer.apiKeySecretRef.key }}

server/server.go

@@ -22,6 +22,8 @@ import (
 	"github.com/kdwils/envoy-proxy-bouncer/template"
 	"google.golang.org/genproto/googleapis/rpc/status"
 	"google.golang.org/grpc"
+	"google.golang.org/grpc/health"
+	"google.golang.org/grpc/health/grpc_health_v1"
 	"google.golang.org/grpc/reflection"
 )
 
@@ -34,9 +36,13 @@ type Server struct {
 	templateStore TemplateStore
 	now           func() time.Time
 	rateLimiter   *RateLimiter
+	healthServer  *health.Server
 }
 
 func NewServer(config config.Config, bouncer Bouncer, captcha Captcha, templateStore TemplateStore, logger *slog.Logger) *Server {
+	healthServer := health.NewServer()
+	healthServer.SetServingStatus("", grpc_health_v1.HealthCheckResponse_NOT_SERVING)
+
 	return &Server{
 		config:        config,
 		bouncer:       bouncer,
@@ -45,6 +51,7 @@ func NewServer(config config.Config, bouncer Bouncer, captcha Captcha, templateS
 		templateStore: templateStore,
 		now:           time.Now,
 		rateLimiter:   NewRateLimiter(10, 20),
+		healthServer:  healthServer,
 	}
 }
 
@@ -100,8 +107,11 @@ func (s *Server) serveGRPC(ctx context.Context, port int) error {
 		grpc.UnaryInterceptor(s.loggerInterceptor),
 	)
 	auth.RegisterAuthorizationServer(grpcServer, s)
+	grpc_health_v1.RegisterHealthServer(grpcServer, s.healthServer)
 	reflection.Register(grpcServer)
 
+	go s.updateHealthStatus(ctx)
+
 	go func() {
 		<-ctx.Done()
 		s.logger.Info("shutting down gRPC server...")
@@ -286,6 +296,34 @@ func (s *Server) handleCaptchaChallenge(w http.ResponseWriter, r *http.Request)
 	w.Write([]byte(html))
 }
 
+func (s *Server) updateHealthStatus(ctx context.Context) {
+	ticker := time.NewTicker(5 * time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ctx.Done():
+			s.healthServer.SetServingStatus("", grpc_health_v1.HealthCheckResponse_NOT_SERVING)
+			return
+		case <-ticker.C:
+			if !s.isReady() {
+				continue
+			}
+
+			s.healthServer.SetServingStatus("", grpc_health_v1.HealthCheckResponse_SERVING)
+			return
+		}
+	}
+}
+
+func (s *Server) isReady() bool {
+	if s.bouncer == nil {
+		return false
+	}
+
+	return s.bouncer.IsReady()
+}
+
 func (s *Server) loggerInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) {
 	reqLogger := slog.New(s.logger.Handler())
 	return handler(logger.WithContext(ctx, reqLogger), req)
@@ -296,6 +334,7 @@ func (s *Server) Check(ctx context.Context, req *auth.CheckRequest) (*auth.Check
 		body, headers := s.renderDeniedResponse(bouncer.NewCheckedRequest("", "", "remediator not initialized", 0, nil, "", nil, nil))
 		return getDeniedResponse(envoy_type.StatusCode_InternalServerError, body, headers), nil
 	}
+
 	result := s.bouncer.Check(ctx, req)
 	s.logger.Info("remediation result", slog.Any("result", result))
 	switch result.Action {

server/services.go

@@ -15,6 +15,7 @@ type Bouncer interface {
 	Sync(ctx context.Context) error
 	Metrics(ctx context.Context) error
 	ExtractRealIPFromHTTP(r *http.Request) string
+	IsReady() bool
 }
 
 //go:generate mockgen -destination=mocks/mock_template_store.go -package=mocks github.com/kdwils/envoy-proxy-bouncer/server TemplateStore