health-check: add storage health check (#5074)

pkg/health/check.go

@@ -15,6 +15,8 @@ const (
 	XDSOther = Check("xds.other")
 	// ZeroBootstrapConfigSave checks whether the Zero bootstrap config was saved
 	ZeroBootstrapConfigSave = Check("zero.bootstrap-config.save")
+	// StorageBackend checks whether the storage backend is healthy
+	StorageBackend = Check("storage.backend")
 )
 
 // ZeroResourceBundle checks whether the Zero resource bundle was applied

pkg/storage/inmemory/backend.go

@@ -20,6 +20,7 @@ import (
 	"github.com/pomerium/pomerium/internal/signal"
 	"github.com/pomerium/pomerium/pkg/cryptutil"
 	"github.com/pomerium/pomerium/pkg/grpc/databroker"
+	"github.com/pomerium/pomerium/pkg/health"
 	"github.com/pomerium/pomerium/pkg/storage"
 )
 
@@ -86,6 +87,9 @@ func New(options ...Option) *Backend {
 			}
 		}()
 	}
+
+	health.ReportOK(health.StorageBackend, health.StrAttr("backend", "in-memory"))
+
 	return backend
 }
 

pkg/storage/postgres/backend.go

@@ -17,6 +17,7 @@ import (
 	"github.com/pomerium/pomerium/internal/signal"
 	"github.com/pomerium/pomerium/pkg/contextutil"
 	"github.com/pomerium/pomerium/pkg/grpc/databroker"
+	"github.com/pomerium/pomerium/pkg/health"
 	"github.com/pomerium/pomerium/pkg/storage"
 )
 
@@ -78,6 +79,16 @@ func New(dsn string, options ...Option) *Backend {
 		return backend.listenForNotifications(ctx)
 	}, time.Millisecond*100)
 
+	go backend.doPeriodically(func(ctx context.Context) error {
+		err := backend.ping(ctx)
+		if err != nil {
+			health.ReportError(health.StorageBackend, err, health.StrAttr("backend", "postgres"))
+		} else {
+			health.ReportOK(health.StorageBackend, health.StrAttr("backend", "postgres"))
+		}
+		return nil
+	}, time.Minute)
+
 	return backend
 }
 
@@ -440,6 +451,15 @@ func (backend *Backend) listenForNotifications(ctx context.Context) error {
 	}
 }
 
+func (backend *Backend) ping(ctx context.Context) error {
+	_, pool, err := backend.init(ctx)
+	if err != nil {
+		return err
+	}
+
+	return pool.Ping(ctx)
+}
+
 // ParseConfig parses a DSN into a pgxpool.Config.
 func ParseConfig(dsn string) (*pgxpool.Config, error) {
 	config, err := pgxpool.ParseConfig(dsn)