87 lines
1.7 KiB
Go
87 lines
1.7 KiB
Go
package healthz
|
|
|
|
import (
|
|
"context"
|
|
|
|
"github.com/cloudwego/hertz/pkg/app"
|
|
"github.com/cloudwego/hertz/pkg/app/server"
|
|
)
|
|
|
|
type Option struct {
|
|
living func() error
|
|
ready func() error
|
|
}
|
|
type healthz struct {
|
|
Option
|
|
}
|
|
|
|
func (h *healthz) SetLiving(living func() error) *healthz {
|
|
h.living = living
|
|
return h
|
|
}
|
|
|
|
func (h *healthz) SetReady(ready func() error) *healthz {
|
|
h.ready = ready
|
|
return h
|
|
}
|
|
|
|
func (h *healthz) Living(ctx context.Context, c *app.RequestContext) {
|
|
if h.living == nil {
|
|
c.JSON(200, map[string]string{"ok": "true"})
|
|
return
|
|
}
|
|
if err := h.living(); err != nil {
|
|
c.JSON(500, map[string]string{"ok": "false"})
|
|
return
|
|
}
|
|
c.JSON(200, map[string]string{"ok": "true"})
|
|
}
|
|
func (h *healthz) Ready(ctx context.Context, c *app.RequestContext) {
|
|
if h.ready == nil {
|
|
c.JSON(200, map[string]string{"ok": "true"})
|
|
return
|
|
}
|
|
if err := h.ready(); err != nil {
|
|
c.JSON(500, map[string]string{"ok": "false"})
|
|
return
|
|
}
|
|
c.JSON(200, map[string]string{"ok": "true"})
|
|
}
|
|
|
|
func (h *healthz) Registry(hertz *server.Hertz) {
|
|
hertz.GET("/living", h.Living)
|
|
hertz.GET("/ready", h.Ready)
|
|
}
|
|
|
|
type Healthz interface {
|
|
SetLiving(living func() error) *healthz
|
|
SetReady(ready func() error) *healthz
|
|
Living(ctx context.Context, c *app.RequestContext)
|
|
Ready(ctx context.Context, c *app.RequestContext)
|
|
Registry(h *server.Hertz)
|
|
}
|
|
|
|
type OptionFunc func(*Option)
|
|
|
|
func NewHealthz(opts ...OptionFunc) Healthz {
|
|
var opt = Option{}
|
|
for _, o := range opts {
|
|
o(&opt)
|
|
}
|
|
return &healthz{
|
|
Option: opt,
|
|
}
|
|
}
|
|
|
|
func WithLiving(living func() error) OptionFunc {
|
|
return func(o *Option) {
|
|
o.living = living
|
|
}
|
|
}
|
|
|
|
func WithReady(ready func() error) OptionFunc {
|
|
return func(o *Option) {
|
|
o.ready = ready
|
|
}
|
|
}
|