Repository: go-kratos/kratos
Pull request #3492 (2024-12-18 09:11:46)

fix(registry): add timeout validation
使用 Consul 做服务发现时，如果 gRPC Client 的 timeout 设置为 0，程序会报错 `discovery create watcher overtime`。

原因是：

1. gRPC Client 会将 `timeout` 传给 discovery；

https://github.com/go-kratos/kratos/blob/861493a20524217061cf329eff0b4e39fb9cbfeb/transport/grpc/client.go#L190-L196

2. discovery 的 `Build` 会因为 `timeout` 为 0，立马超时，得到 `ErrWatcherCreateTimeout` 的错误。

https://github.com/go-kratos/kratos/blob/861493a20524217061cf329eff0b4e39fb9cbfeb/transport/grpc/resolver/discovery/builder.go#L97-L102


shenqidebaozi: This pull request includes several changes to improve the handling of timeouts and error checking in service resolution and gRPC resolver building. The most important changes include modifications to the `resolve` method in `contrib/registry/consul/registry.go` and the `Build` method in `transport/grpc/resolver/discovery/builder.go`.

Improvements to service resolution:

* [`contrib/registry/consul/registry.go`](diffhunk://#diff-04cd5872cae1887118fcbff9cccb290f59a0d99c8f17f981bfab2f88e00784ccR218-R228): Added a conditional timeout to the `resolve` method, which now uses a new `listServices` function to handle service listing with or without a timeout. [[1]](diffhunk://#diff-04cd5872cae1887118fcbff9cccb290f59a0d99c8f17f981bfab2f88e00784ccR218-R228) [[2]](diffhunk://#diff-04cd5872cae1887118fcbff9cccb290f59a0d99c8f17f981bfab2f88e00784ccL235-R242)

Enhancements to gRPC resolver building:

* [`transport/grpc/resolver/discovery/builder.go`](diffhunk://#diff-c4f4b1c1479a6baf4c5cf886c986ba13c2d1a62286f056b372910aebedce1157R97-R107): Modified the `Build` method to include a conditional timeout check, ensuring that the method properly handles the timeout scenario and waits for the operation to complete if no timeout is set.

Error handling improvements in tests:

* [`transport/grpc/resolver/discovery/builder_test.go`](diffhunk://#diff-ed7bb15092064eebecc40b384deb40e6411711e7f1ab1349ed6e628540d56d3cL110-R111): Updated the `TestBuilder_Build` test to correctly check for the absence of an error, ensuring the test's expected behavior aligns with the new implementation.

Patch:
diff --git a/contrib/registry/consul/registry.go b/contrib/registry/consul/registry.go
index 5b098384..42fa4a37 100644
--- a/contrib/registry/consul/registry.go
+++ b/contrib/registry/consul/registry.go
@@ -215,10 +215,17 @@ func (r *Registry) Watch(ctx context.Context, name string) (registry.Watcher, er
 }
 
 func (r *Registry) resolve(ctx context.Context, ss *serviceSet) error {
-	timeoutCtx, cancel := context.WithTimeout(ctx, r.timeout)
-	defer cancel()
+	listServices := r.cli.Service
+	if r.timeout > 0 {
+		listServices = func(ctx context.Context, service string, index uint64, passingOnly bool) ([]*registry.ServiceInstance, uint64, error) {
+			timeoutCtx, cancel := context.WithTimeout(ctx, r.timeout)
+			defer cancel()
 
-	services, idx, err := r.cli.Service(timeoutCtx, ss.serviceName, 0, true)
+			return r.cli.Service(timeoutCtx, service, index, passingOnly)
+		}
+	}
+
+	services, idx, err := listServices(ctx, ss.serviceName, 0, true)
 	if err != nil {
 		return err
 	}
@@ -232,9 +239,7 @@ func (r *Registry) resolve(ctx context.Context, ss *serviceSet) error {
 		for {
 			select {
 			case <-ticker.C:
-				timeoutCtx, cancel := context.WithTimeout(context.Background(), r.timeout)
-				tmpService, tmpIdx, err := r.cli.Service(timeoutCtx, ss.serviceName, idx, true)
-				cancel()
+				tmpService, tmpIdx, err := listServices(context.Background(), ss.serviceName, idx, true)
 				if err != nil {
 					time.Sleep(time.Second)
 					continue
diff --git a/transport/grpc/resolver/discovery/builder.go b/transport/grpc/resolver/discovery/builder.go
index e8e98726..dfed56bf 100644
--- a/transport/grpc/resolver/discovery/builder.go
+++ b/transport/grpc/resolver/discovery/builder.go
@@ -94,11 +94,16 @@ func (b *builder) Build(target resolver.Target, cc resolver.ClientConn, _ resolv
 	}()
 
 	var err error
-	select {
-	case <-done:
+	if b.timeout > 0 {
+		select {
+		case <-done:
+			err = watchRes.err
+		case <-time.After(b.timeout):
+			err = ErrWatcherCreateTimeout
+		}
+	} else {
+		<-done
 		err = watchRes.err
-	case <-time.After(b.timeout):
-		err = ErrWatcherCreateTimeout
 	}
 	if err != nil {
 		cancel()


Test patch:
diff --git a/transport/grpc/resolver/discovery/builder_test.go b/transport/grpc/resolver/discovery/builder_test.go
index 0f27078a..1c855830 100644
--- a/transport/grpc/resolver/discovery/builder_test.go
+++ b/transport/grpc/resolver/discovery/builder_test.go
@@ -107,7 +107,7 @@ func TestBuilder_Build(t *testing.T) {
 		&mockConn{},
 		resolver.BuildOptions{},
 	)
-	if err == nil {
-		t.Errorf("expected error, got %v", err)
+	if err != nil {
+		t.Errorf("expected no error, got %v", err)
 	}
 }