From b354f185c04421bd81c8890f89f2079f1f976691 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=8C=85=E5=AD=90?= <baozhecheng@foxmail.com>
Date: Fri, 26 Aug 2022 10:54:02 +0800
Subject: [PATCH] feat(registry): consul support user custom checks (#2317)

---
 contrib/registry/consul/client.go   | 5 +++++
 contrib/registry/consul/registry.go | 9 +++++++++
 2 files changed, 14 insertions(+)

diff --git a/contrib/registry/consul/client.go b/contrib/registry/consul/client.go
index e97ec7e42..1b39975a0 100644
--- a/contrib/registry/consul/client.go
+++ b/contrib/registry/consul/client.go
@@ -29,6 +29,8 @@ type Client struct {
 	heartbeat bool
 	// deregisterCriticalServiceAfter time interval in seconds
 	deregisterCriticalServiceAfter int
+	// serviceChecks  user custom checks
+	serviceChecks api.AgentServiceChecks
 }
 
 // NewClient creates consul client
@@ -139,6 +141,9 @@ func (c *Client) Register(_ context.Context, svc *registry.ServiceInstance, enab
 		})
 	}
 
+	// custom checks
+	asr.Checks = append(asr.Checks, c.serviceChecks...)
+
 	err := c.cli.Agent().ServiceRegister(asr)
 	if err != nil {
 		return err
diff --git a/contrib/registry/consul/registry.go b/contrib/registry/consul/registry.go
index 9da95d77b..d1c47d774 100644
--- a/contrib/registry/consul/registry.go
+++ b/contrib/registry/consul/registry.go
@@ -63,6 +63,15 @@ func WithDeregisterCriticalServiceAfter(interval int) Option {
 	}
 }
 
+// WithServiceCheck with service checks
+func WithServiceCheck(checks ...*api.AgentServiceCheck) Option {
+	return func(o *Registry) {
+		if o.cli != nil {
+			o.cli.serviceChecks = checks
+		}
+	}
+}
+
 // Config is consul registry config
 type Config struct {
 	*api.Config