Repository: go-kratos/kratos
Pull request #1738 (2021-12-31 12:42:53)

feat(registry/consul): add WithHeartbeat option
<!--
🎉 Thanks for sending a pull request to Kratos! Here are some tips for you:

1. If this is your first time contributing to Kratos, please read our contribution guide: https://go-kratos.dev/en/docs/community/contribution/
2. Ensure you have added or ran the appropriate tests and lint for your PR, please use `make lint` and `make test` before filing your PR, use `make clean` to tidy your go mod.
3. If the PR is unfinished, you may need mark it as a WIP(Work In Progress) PR or draft PR
4. Please use a conventional commits format title: `<type>[optional scope]: <description>`
    
    Some suggestion on <type>:

    fix: A bug fix
    feat: A new feature
    test: Adding missing tests or correcting existing tests
    refactor: A code change that neither fixes a bug nor adds a feature
    break: Changes has break change

    docs: Documentation only changes
    deps: Changes external dependencies
    style: Changes that do not affect the meaning of the code (white-space, formatting, etc)
    chore Daily work, examples, etc.
    ci: Changes to our CI configuration files and scripts
-->

#### Description (what this PR does / why we need it):
<!--
* The description should include the motivation for this PR or contrast this with previous behavior
-->
Add WithHeartbeat option, so user can choose healthcheck or heartbeat, or both.

#### Which issue(s) this PR fixes (resolves / be part of):
<!--
* Automatically closes linked issue when PR is merged.
* If you PR is not fully resolved issue, please use `part of #<issue number>` instead.

Usage: `fixes/resolves #<issue number>`, or `fixes/resolves (paste link of issue)`.
-->
fixes #

#### Other special notes for reviewer:
<!--
* Somethings that need extra attention for reviewer
* Some additional notes, TODO list, etc.
-->


codecov-commenter: # [Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1738?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) Report
> Merging [#1738](https://codecov.io/gh/go-kratos/kratos/pull/1738?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (f2006f8) into [main](https://codecov.io/gh/go-kratos/kratos/commit/17201cd2845677643c77db8a1d5ca1228314ec2f?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (17201cd) will **increase** coverage by `0.05%`.
> The diff coverage is `n/a`.

[![Impacted file tree graph](https://codecov.io/gh/go-kratos/kratos/pull/1738/graphs/tree.svg?width=650&height=150&src=pr&token=XwnmBTz78x&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)](https://codecov.io/gh/go-kratos/kratos/pull/1738?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)

```diff
@@            Coverage Diff             @@
##             main    #1738      +/-   ##
==========================================
+ Coverage   81.21%   81.27%   +0.05%     
==========================================
  Files          78       78              
  Lines        3476     3476              
==========================================
+ Hits         2823     2825       +2     
+ Misses        448      446       -2     
  Partials      205      205              
```


| [Impacted Files](https://codecov.io/gh/go-kratos/kratos/pull/1738?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) | Coverage Δ | |
|---|---|---|
| [internal/context/context.go](https://codecov.io/gh/go-kratos/kratos/pull/1738/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-aW50ZXJuYWwvY29udGV4dC9jb250ZXh0Lmdv) | `100.00% <0.00%> (+2.98%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1738?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1738?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos). Last update [17201cd...f2006f8](https://codecov.io/gh/go-kratos/kratos/pull/1738?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos).


Patch:
diff --git a/contrib/registry/consul/client.go b/contrib/registry/consul/client.go
index 62654db0..b5e19a17 100644
--- a/contrib/registry/consul/client.go
+++ b/contrib/registry/consul/client.go
@@ -23,6 +23,8 @@ type Client struct {
 	resolver ServiceResolver
 	// healthcheck time interval in seconds
 	healthcheckInterval int
+	// heartbeat enable heartbeat
+	heartbeat bool
 }
 
 // NewClient creates consul client
@@ -31,6 +33,7 @@ func NewClient(cli *api.Client) *Client {
 		cli:                 cli,
 		resolver:            defaultResolver,
 		healthcheckInterval: 10,
+		heartbeat:           true,
 	}
 	c.ctx, c.cancel = context.WithCancel(context.Background())
 	return c
@@ -123,18 +126,20 @@ func (c *Client) Register(_ context.Context, svc *registry.ServiceInstance, enab
 		return err
 	}
 	_ = c.cli.Agent().UpdateTTL("service:"+svc.ID, "pass", "pass")
-	go func() {
-		ticker := time.NewTicker(time.Second * time.Duration(c.healthcheckInterval))
-		defer ticker.Stop()
-		for {
-			select {
-			case <-ticker.C:
-				_ = c.cli.Agent().UpdateTTL("service:"+svc.ID, "pass", "pass")
-			case <-c.ctx.Done():
-				return
+	if c.heartbeat {
+		go func() {
+			ticker := time.NewTicker(time.Second * time.Duration(c.healthcheckInterval))
+			defer ticker.Stop()
+			for {
+				select {
+				case <-ticker.C:
+					_ = c.cli.Agent().UpdateTTL("service:"+svc.ID, "pass", "pass")
+				case <-c.ctx.Done():
+					return
+				}
 			}
-		}
-	}()
+		}()
+	}
 	return nil
 }
 
diff --git a/contrib/registry/consul/registry.go b/contrib/registry/consul/registry.go
index e2dd7a6e..36b64859 100644
--- a/contrib/registry/consul/registry.go
+++ b/contrib/registry/consul/registry.go
@@ -26,6 +26,15 @@ func WithHealthCheck(enable bool) Option {
 	}
 }
 
+// WithHeartbeat enable or disable heartbeat
+func WithHeartbeat(enable bool) Option {
+	return func(o *Registry) {
+		if o.cli != nil {
+			o.cli.heartbeat = enable
+		}
+	}
+}
+
 // WithServiceResolver with endpoint function option.
 func WithServiceResolver(fn ServiceResolver) Option {
 	return func(o *Registry) {
@@ -144,33 +153,42 @@ func (r *Registry) Watch(ctx context.Context, name string) (registry.Watcher, er
 	}
 
 	if !ok {
-		go r.resolve(set)
+		err := r.resolve(set)
+		if err != nil {
+			return nil, err
+		}
 	}
 	return w, nil
 }
 
-func (r *Registry) resolve(ss *serviceSet) {
+func (r *Registry) resolve(ss *serviceSet) error {
 	ctx, cancel := context.WithTimeout(context.Background(), time.Second*10)
 	services, idx, err := r.cli.Service(ctx, ss.serviceName, 0, true)
 	cancel()
-	if err == nil && len(services) > 0 {
+	if err != nil {
+		return err
+	} else if len(services) > 0 {
 		ss.broadcast(services)
 	}
-	ticker := time.NewTicker(time.Second)
-	defer ticker.Stop()
-	for {
-		<-ticker.C
-		ctx, cancel := context.WithTimeout(context.Background(), time.Second*120)
-		tmpService, tmpIdx, err := r.cli.Service(ctx, ss.serviceName, idx, true)
-		cancel()
-		if err != nil {
-			time.Sleep(time.Second)
-			continue
+	go func() {
+		ticker := time.NewTicker(time.Second)
+		defer ticker.Stop()
+		for {
+			<-ticker.C
+			ctx, cancel := context.WithTimeout(context.Background(), time.Second*120)
+			tmpService, tmpIdx, err := r.cli.Service(ctx, ss.serviceName, idx, true)
+			cancel()
+			if err != nil {
+				time.Sleep(time.Second)
+				continue
+			}
+			if len(tmpService) != 0 && tmpIdx != idx {
+				services = tmpService
+				ss.broadcast(services)
+			}
+			idx = tmpIdx
 		}
-		if len(tmpService) != 0 && tmpIdx != idx {
-			services = tmpService
-			ss.broadcast(services)
-		}
-		idx = tmpIdx
-	}
+	}()
+
+	return nil
 }


Test patch:
diff --git a/contrib/registry/consul/registry_test.go b/contrib/registry/consul/registry_test.go
index c95f6b02..c0f9b4ab 100644
--- a/contrib/registry/consul/registry_test.go
+++ b/contrib/registry/consul/registry_test.go
@@ -25,7 +25,7 @@ func tcpServer(t *testing.T, lis net.Listener) {
 }
 
 func TestRegister(t *testing.T) {
-	addr := fmt.Sprintf("%s:8081", getIntranetIP())
+	addr := fmt.Sprintf("%s:9091", getIntranetIP())
 	lis, err := net.Listen("tcp", addr)
 	if err != nil {
 		t.Errorf("listen tcp %s failed!", addr)
@@ -38,7 +38,12 @@ func TestRegister(t *testing.T) {
 	if err != nil {
 		t.Fatalf("create consul client failed: %v", err)
 	}
-	r := New(cli)
+	opts := []Option{
+		WithHeartbeat(true),
+		WithHealthCheck(true),
+		WithHealthCheckInterval(5),
+	}
+	r := New(cli, opts...)
 	assert.Nil(t, err)
 	version := strconv.FormatInt(time.Now().Unix(), 10)
 	svc := &registry.ServiceInstance{