From 27f10ca798c2e19cd0f1f36486837472ea57f6ae Mon Sep 17 00:00:00 2001
From: timerzz <zhhgzhhgzhhg@163.com>
Date: Mon, 26 Aug 2024 16:28:43 +0800
Subject: [PATCH] =?UTF-8?q?feat=20=E6=B7=BB=E5=8A=A0=E6=8E=A2=E9=92=88?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 pkg/web/probe.go    | 56 +++++++++++++++++++++++++++++++++++++++++++++
 pkg/web/register.go |  7 ++++++
 2 files changed, 63 insertions(+)
 create mode 100644 pkg/web/probe.go
 create mode 100644 pkg/web/register.go

diff --git a/pkg/web/probe.go b/pkg/web/probe.go
new file mode 100644
index 0000000..8cfd63f
--- /dev/null
+++ b/pkg/web/probe.go
@@ -0,0 +1,56 @@
+package web
+
+import "github.com/gofiber/fiber/v3"
+
+// 注册健康探针
+type Probe struct {
+	// 存活探针
+	livingHook func(ctx fiber.Ctx) error
+	// 就绪探针
+	readyHook func(ctx fiber.Ctx) error
+	// 启动探针
+	startHook func(ctx fiber.Ctx) error
+}
+
+func NewProbe() *Probe {
+	return &Probe{}
+}
+
+// 设置存活探针
+func (p *Probe) SetLivingProbe(hook func(ctx fiber.Ctx) error) *Probe {
+	p.livingHook = hook
+	return p
+}
+
+// 设置就绪探针
+func (p *Probe) SetReadyProbe(hook func(ctx fiber.Ctx) error) *Probe {
+	p.readyHook = hook
+	return p
+}
+
+// 设置启动探针
+func (p *Probe) SetStartProbe(hook func(ctx fiber.Ctx) error) *Probe {
+	p.startHook = hook
+	return p
+}
+
+// 默认探针
+func defaultProbe(ctx fiber.Ctx) error {
+	return ctx.SendStatus(fiber.StatusOK)
+}
+
+func (p *Probe) Registry(r fiber.Router) {
+	if p.livingHook == nil {
+		p.SetLivingProbe(defaultProbe)
+	}
+	if p.readyHook == nil {
+		p.SetReadyProbe(defaultProbe)
+	}
+	if p.startHook == nil {
+		p.SetStartProbe(defaultProbe)
+	}
+
+	r.Get("/ready", p.readyHook)
+	r.Get("/living", p.livingHook)
+	r.Get("/start", p.startHook)
+}
diff --git a/pkg/web/register.go b/pkg/web/register.go
new file mode 100644
index 0000000..e00862d
--- /dev/null
+++ b/pkg/web/register.go
@@ -0,0 +1,7 @@
+package web
+
+import "github.com/gofiber/fiber/v3"
+
+type Register interface {
+	Registry(r fiber.Router)
+}