diff --git a/client/clients/router/client.go b/client/clients/router/client.go index b7ab53303fb..e59aa55549c 100644 --- a/client/clients/router/client.go +++ b/client/clients/router/client.go @@ -42,6 +42,7 @@ import ( "github.com/tikv/pd/client/pkg/batch" cctx "github.com/tikv/pd/client/pkg/connectionctx" "github.com/tikv/pd/client/pkg/retry" + "github.com/tikv/pd/client/pkg/utils/grpcutil" sd "github.com/tikv/pd/client/servicediscovery" ) @@ -842,6 +843,10 @@ func (c *Cli) processRequestsInner( queryReq := buildQueryRegionRequest(c.svcDiscovery.GetClusterID(), requests) start := time.Now() if err := send(queryReq); err != nil { + err = grpcutil.ResolveStreamSendError(err, func() error { + _, err := recv() + return err + }) metrics.RequestFailedDurationQueryRegion.Observe(time.Since(start).Seconds()) return nil, err } diff --git a/client/clients/router/expected_service_test.go b/client/clients/router/expected_service_test.go new file mode 100644 index 00000000000..fd6bd14bfe7 --- /dev/null +++ b/client/clients/router/expected_service_test.go @@ -0,0 +1,72 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package router + +import ( + "context" + "net" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/connectivity" + "google.golang.org/grpc/health" + healthpb "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" + + "github.com/pingcap/kvproto/pkg/pdpb" + + "github.com/tikv/pd/client/pkg/batch" + "github.com/tikv/pd/client/pkg/utils/grpcutil" + sd "github.com/tikv/pd/client/servicediscovery" +) + +// TestProcessRequestsClosesConnOnRejectedStream covers a pre-opened QueryRegion +// stream rejected by an endpoint which does not provide the PD service: Send +// returns io.EOF, and the connection must still be closed to be recreated. +func TestProcessRequestsClosesConnOnRejectedStream(t *testing.T) { + re := require.New(t) + lis, err := net.Listen("tcp", "127.0.0.1:0") + re.NoError(err) + srv := grpc.NewServer() + healthpb.RegisterHealthServer(srv, health.NewServer()) + go func() { + _ = srv.Serve(lis) + }() + defer srv.Stop() + + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + cc, err := grpcutil.GetClientConn(ctx, "http://"+lis.Addr().String(), nil, + grpcutil.ExpectedServiceDialOptions("pdpb.PD")...) + re.NoError(err) + defer cc.Close() + stream, err := pdpb.NewPDClient(cc).QueryRegion(ctx) + re.NoError(err) + // Wait for the stream to be rejected. + re.Eventually(func() bool { + return stream.Send(&pdpb.QueryRegionRequest{}) != nil + }, 5*time.Second, 10*time.Millisecond) + + c := &Cli{ + svcDiscovery: sd.NewMockServiceDiscovery(nil, nil), + batchController: batch.NewController[*Request](1, requestFinisher(nil, nil), nil), + } + _, err = c.processRequestsInner(stream.Send, stream.Recv, false, false) + re.Equal(codes.Unimplemented, status.Code(err)) + re.Equal(connectivity.Shutdown, cc.GetState()) +} diff --git a/client/pkg/utils/grpcutil/expected_service_test.go b/client/pkg/utils/grpcutil/expected_service_test.go new file mode 100644 index 00000000000..1cd1ae7f6e7 --- /dev/null +++ b/client/pkg/utils/grpcutil/expected_service_test.go @@ -0,0 +1,227 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package grpcutil + +import ( + "context" + "io" + "net" + "sync" + "testing" + "time" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/connectivity" + "google.golang.org/grpc/health" + healthpb "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" + + "github.com/pingcap/kvproto/pkg/keyspacepb" + "github.com/pingcap/kvproto/pkg/pdpb" +) + +// startGRPCServer starts a gRPC server which only registers the health service, +// simulating an endpoint that is owned by another kind of service. +func startGRPCServer(t *testing.T, register func(*grpc.Server)) string { + lis, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + srv := grpc.NewServer() + healthpb.RegisterHealthServer(srv, health.NewServer()) + if register != nil { + register(srv) + } + go func() { + _ = srv.Serve(lis) + }() + t.Cleanup(srv.Stop) + return "http://" + lis.Addr().String() +} + +func dialWithExpectedService(t *testing.T, url, service string) *grpc.ClientConn { + ctx, cancel := context.WithTimeout(context.Background(), dialTimeout) + defer cancel() + cc, err := GetClientConn(ctx, url, nil, ExpectedServiceDialOptions(service)...) + require.NoError(t, err) + t.Cleanup(func() { + _ = cc.Close() + }) + return cc +} + +func TestExpectedServiceClosesConnOnUnknownService(t *testing.T) { + re := require.New(t) + url := startGRPCServer(t, nil) + cc := dialWithExpectedService(t, url, "pdpb.PD") + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _, err := pdpb.NewPDClient(cc).GetMembers(ctx, &pdpb.GetMembersRequest{}) + re.Equal(codes.Unimplemented, status.Code(err)) + re.Equal(connectivity.Shutdown, cc.GetState()) +} + +func TestExpectedServiceClosesConnOnUnknownServiceStream(t *testing.T) { + re := require.New(t) + url := startGRPCServer(t, nil) + cc := dialWithExpectedService(t, url, "pdpb.PD") + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + stream, err := pdpb.NewPDClient(cc).Tso(ctx) + re.NoError(err) + _, err = stream.Recv() + re.Equal(codes.Unimplemented, status.Code(err)) + re.Equal(connectivity.Shutdown, cc.GetState()) +} + +// TestExpectedServiceClosesConnOnStreamSendEOF covers a pre-opened stream: once +// the server rejects it, Send returns io.EOF and only Recv reports the status. +func TestExpectedServiceClosesConnOnStreamSendEOF(t *testing.T) { + re := require.New(t) + url := startGRPCServer(t, nil) + cc := dialWithExpectedService(t, url, "pdpb.PD") + + stream, err := pdpb.NewPDClient(cc).QueryRegion(t.Context()) + re.NoError(err) + var sendErr error + re.Eventually(func() bool { + sendErr = stream.Send(&pdpb.QueryRegionRequest{}) + return sendErr != nil + }, 5*time.Second, 10*time.Millisecond) + re.ErrorIs(sendErr, io.EOF) + re.NotEqual(connectivity.Shutdown, cc.GetState()) + + err = ResolveStreamSendError(sendErr, func() error { + _, err := stream.Recv() + return err + }) + re.Equal(codes.Unimplemented, status.Code(err)) + re.Equal(connectivity.Shutdown, cc.GetState()) +} + +func TestResolveStreamSendError(t *testing.T) { + re := require.New(t) + recvErr := status.Error(codes.Unimplemented, "unknown service pdpb.PD") + recvCalled := false + recv := func(err error) func() error { + return func() error { + recvCalled = true + return err + } + } + + re.NoError(ResolveStreamSendError(nil, recv(recvErr))) + re.False(recvCalled) + otherErr := status.Error(codes.Unavailable, "unavailable") + re.Equal(otherErr, ResolveStreamSendError(otherErr, recv(recvErr))) + re.False(recvCalled) + re.Equal(recvErr, ResolveStreamSendError(io.EOF, recv(recvErr))) + re.True(recvCalled) + // Keep io.EOF if the stream status can not be got by receiving. + re.Equal(io.EOF, ResolveStreamSendError(io.EOF, recv(nil))) +} + +func TestExpectedServiceKeepsConnOnUnknownMethod(t *testing.T) { + re := require.New(t) + url := startGRPCServer(t, nil) + cc := dialWithExpectedService(t, url, healthpb.Health_ServiceDesc.ServiceName) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + err := cc.Invoke(ctx, "/grpc.health.v1.Health/NoSuchMethod", + &healthpb.HealthCheckRequest{}, &healthpb.HealthCheckResponse{}) + re.Equal(codes.Unimplemented, status.Code(err)) + re.NotEqual(connectivity.Shutdown, cc.GetState()) +} + +func TestExpectedServiceKeepsConnOnOtherUnknownService(t *testing.T) { + re := require.New(t) + // The endpoint is a real PD, but it does not provide the keyspace service. + url := startGRPCServer(t, func(s *grpc.Server) { + pdpb.RegisterPDServer(s, &pdpb.UnimplementedPDServer{}) + }) + cc := dialWithExpectedService(t, url, "pdpb.PD") + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + _, err := keyspacepb.NewKeyspaceClient(cc).LoadKeyspace(ctx, &keyspacepb.LoadKeyspaceRequest{}) + re.Equal(codes.Unimplemented, status.Code(err)) + re.NotEqual(connectivity.Shutdown, cc.GetState()) + // An unimplemented method of the expected service must not close the connection either. + _, err = pdpb.NewPDClient(cc).GetMembers(ctx, &pdpb.GetMembersRequest{}) + re.Equal(codes.Unimplemented, status.Code(err)) + re.NotEqual(connectivity.Shutdown, cc.GetState()) +} + +func TestGetOrCreateGRPCConnRecreatesShutdownConn(t *testing.T) { + re := require.New(t) + url := startGRPCServer(t, nil) + clientConns := &sync.Map{} + + cc, err := GetOrCreateGRPCConn(context.Background(), clientConns, url, nil) + re.NoError(err) + re.NoError(cc.Close()) + + newCC, err := GetOrCreateGRPCConn(context.Background(), clientConns, url, nil) + re.NoError(err) + defer newCC.Close() + re.NotSame(cc, newCC) + re.NotEqual(connectivity.Shutdown, newCC.GetState()) + cached, ok := clientConns.Load(url) + re.True(ok) + re.Same(newCC, cached) +} + +// TestUnknownServiceErrorFormat pins the message returned by the gRPC server for +// an unregistered service, which isUnknownServiceError relies on. If it fails +// after upgrading gRPC, update unknownServiceErrorPrefix accordingly, otherwise +// the connections reaching another service will never be closed. +func TestUnknownServiceErrorFormat(t *testing.T) { + re := require.New(t) + url := startGRPCServer(t, nil) + ctx, cancel := context.WithTimeout(context.Background(), dialTimeout) + defer cancel() + cc, err := GetClientConn(ctx, url, nil) + re.NoError(err) + defer cc.Close() + + const formatChangedMsg = "gRPC changed the unknown service error format, update unknownServiceErrorPrefix" + checkUnknownServiceError := func(err error) { + s, ok := status.FromError(err) + re.True(ok, formatChangedMsg) + re.Equal(codes.Unimplemented, s.Code(), formatChangedMsg) + re.Equal(unknownServiceErrorPrefix+"pdpb.PD", s.Message(), formatChangedMsg) + re.True(isUnknownServiceError(err, "pdpb.PD")) + re.False(isUnknownServiceError(err, "tsopb.TSO")) + } + + callCtx, callCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer callCancel() + _, err = pdpb.NewPDClient(cc).GetMembers(callCtx, &pdpb.GetMembersRequest{}) + checkUnknownServiceError(err) + + stream, err := pdpb.NewPDClient(cc).Tso(callCtx) + re.NoError(err) + _, err = stream.Recv() + checkUnknownServiceError(err) + + // An unknown method of a registered service must not be treated as an unknown service. + err = cc.Invoke(callCtx, "/grpc.health.v1.Health/NoSuchMethod", + &healthpb.HealthCheckRequest{}, &healthpb.HealthCheckResponse{}) + re.Equal(codes.Unimplemented, status.Code(err)) + re.False(isUnknownServiceError(err, healthpb.Health_ServiceDesc.ServiceName)) +} diff --git a/client/pkg/utils/grpcutil/grpcutil.go b/client/pkg/utils/grpcutil/grpcutil.go index a5a43cb7ac1..f81c17426ab 100644 --- a/client/pkg/utils/grpcutil/grpcutil.go +++ b/client/pkg/utils/grpcutil/grpcutil.go @@ -17,6 +17,7 @@ package grpcutil import ( "context" "crypto/tls" + "io" "net" "net/url" "sync" @@ -26,6 +27,7 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/backoff" "google.golang.org/grpc/codes" + "google.golang.org/grpc/connectivity" "google.golang.org/grpc/credentials" "google.golang.org/grpc/credentials/insecure" "google.golang.org/grpc/metadata" @@ -200,8 +202,13 @@ func getValueFromMetadata(ctx context.Context, key string, f func(context.Contex func GetOrCreateGRPCConn(ctx context.Context, clientConns *sync.Map, url string, tlsCfg *tls.Config, opt ...grpc.DialOption) (*grpc.ClientConn, error) { conn, ok := clientConns.Load(url) if ok { - // TODO: check the connection state. - return conn.(*grpc.ClientConn), nil + cc := conn.(*grpc.ClientConn) + if cc.GetState() != connectivity.Shutdown { + return cc, nil + } + // The connection has been closed, e.g., it reached an endpoint without the + // expected service, so drop it and create a new one. + clientConns.CompareAndDelete(url, conn) } dCtx, cancel := context.WithTimeout(ctx, dialTimeout) defer cancel() @@ -248,3 +255,99 @@ func GetCalleeID(addr string) string { } return callerID } + +// ExpectedServiceDialOptions returns the dial options which close the connection +// once the remote endpoint reports that it does not provide the expected service. +// It usually means the address has been reused by another kind of service, e.g., +// the stale DNS record points to a recycled Pod IP. Closing the connection lets +// the caller drop it and reconnect instead of reusing it forever. +// +// Only the expected service is checked, because a valid endpoint may not provide +// other services or methods, e.g., an older server or a different deployment mode. +func ExpectedServiceDialOptions(service string) []grpc.DialOption { + return []grpc.DialOption{ + grpc.WithChainUnaryInterceptor(unaryExpectedServiceInterceptor(service)), + grpc.WithChainStreamInterceptor(streamExpectedServiceInterceptor(service)), + } +} + +func unaryExpectedServiceInterceptor(service string) grpc.UnaryClientInterceptor { + return func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error { + err := invoker(ctx, method, req, reply, cc, opts...) + closeIfServiceMissing(cc, service, err) + return err + } +} + +func streamExpectedServiceInterceptor(service string) grpc.StreamClientInterceptor { + return func(ctx context.Context, desc *grpc.StreamDesc, cc *grpc.ClientConn, method string, streamer grpc.Streamer, opts ...grpc.CallOption) (grpc.ClientStream, error) { + stream, err := streamer(ctx, desc, cc, method, opts...) + if err != nil { + closeIfServiceMissing(cc, service, err) + return nil, err + } + return &expectedServiceClientStream{ClientStream: stream, cc: cc, service: service}, nil + } +} + +// expectedServiceClientStream checks the received error, since the stream is +// created lazily and the server returns the unknown service error on receiving. +type expectedServiceClientStream struct { + grpc.ClientStream + cc *grpc.ClientConn + service string +} + +// RecvMsg implements grpc.ClientStream. +func (s *expectedServiceClientStream) RecvMsg(m any) error { + err := s.ClientStream.RecvMsg(m) + closeIfServiceMissing(s.cc, s.service, err) + return err +} + +func closeIfServiceMissing(cc *grpc.ClientConn, service string, err error) { + // Concurrent RPCs on the same connection may all fail, close it only once. + if !isUnknownServiceError(err, service) || cc.GetState() == connectivity.Shutdown { + return + } + log.Warn("[pd] close the grpc connection since the remote endpoint does not provide the expected service", + zap.String("target", cc.Target()), zap.String("service", service), errs.ZapError(err)) + if closeErr := cc.Close(); closeErr != nil { + log.Warn("[pd] failed to close grpc clientConn", errs.ZapError(errs.ErrCloseGRPCConn, closeErr)) + } +} + +// ResolveStreamSendError returns the status of the stream if sending fails with +// io.EOF, since gRPC only reports why the stream is terminated, e.g., it is +// rejected for an unknown service, on receiving. Receiving also makes the +// expected service check observe the status. recv is only called on io.EOF. +func ResolveStreamSendError(sendErr error, recv func() error) error { + if sendErr != io.EOF { + return sendErr + } + if err := recv(); err != nil { + return err + } + return sendErr +} + +// unknownServiceErrorPrefix is the prefix of the message returned by the gRPC +// server when the requested service is not registered, see +// `grpc.Server.handleStream`. gRPC does not expose a typed error for it, and +// `codes.Unimplemented` is also returned for an unknown method of a registered +// service, so the message has to be matched. TestUnknownServiceErrorFormat pins +// the format to catch any change when upgrading gRPC. +const unknownServiceErrorPrefix = "unknown service " + +// isUnknownServiceError returns true if the error indicates that the remote +// endpoint does not provide the given gRPC service. +func isUnknownServiceError(err error, service string) bool { + if err == nil { + return false + } + s, ok := status.FromError(err) + if !ok || s.Code() != codes.Unimplemented { + return false + } + return s.Message() == unknownServiceErrorPrefix+service +} diff --git a/client/resource_manager_client.go b/client/resource_manager_client.go index 0fb5cf87d85..15feb397cac 100644 --- a/client/resource_manager_client.go +++ b/client/resource_manager_client.go @@ -30,6 +30,7 @@ import ( "github.com/tikv/pd/client/constants" "github.com/tikv/pd/client/errs" "github.com/tikv/pd/client/opt" + "github.com/tikv/pd/client/pkg/utils/grpcutil" ) type actionType int @@ -424,12 +425,16 @@ tokenRequestLoop: func (c *innerClient) processTokenRequests(stream rmpb.ResourceManager_AcquireTokenBucketsClient, t *tokenRequest) error { req := t.Request - if err := stream.Send(req); err != nil { - err = errors.WithStack(err) - t.done <- err - return err + var resp *rmpb.TokenBucketsResponse + err := stream.Send(req) + if err == nil { + resp, err = stream.Recv() + } else { + err = grpcutil.ResolveStreamSendError(err, func() error { + _, err := stream.Recv() + return err + }) } - resp, err := stream.Recv() if err != nil { c.gRPCErrorHandler(err) c.resourceManagerErrorHandler(err) diff --git a/client/resource_manager_client_test.go b/client/resource_manager_client_test.go index d295989c416..a8549971223 100644 --- a/client/resource_manager_client_test.go +++ b/client/resource_manager_client_test.go @@ -28,6 +28,9 @@ import ( "github.com/stretchr/testify/require" "google.golang.org/grpc" "google.golang.org/grpc/codes" + "google.golang.org/grpc/connectivity" + "google.golang.org/grpc/health" + healthpb "google.golang.org/grpc/health/grpc_health_v1" "google.golang.org/grpc/status" "github.com/pingcap/kvproto/pkg/meta_storagepb" @@ -767,3 +770,40 @@ func TestLeaderChangeDuringServiceModeSwitch(t *testing.T) { require.Equal(t, pdpb.ServiceMode_PD_SVC_MODE, inner.serviceMode) inner.RUnlock() } + +// TestProcessTokenRequestsClosesConnOnRejectedStream covers a token stream +// rejected by an endpoint which does not provide the resource manager service: +// Send returns io.EOF, and the connection must still be closed to be recreated. +func TestProcessTokenRequestsClosesConnOnRejectedStream(t *testing.T) { + re := require.New(t) + lis, err := net.Listen("tcp", "127.0.0.1:0") + re.NoError(err) + srv := grpc.NewServer() + healthpb.RegisterHealthServer(srv, health.NewServer()) + go func() { + _ = srv.Serve(lis) + }() + defer srv.Stop() + + ctx, cancel := context.WithTimeout(t.Context(), 5*time.Second) + defer cancel() + cc, err := grpcutil.GetClientConn(ctx, "http://"+lis.Addr().String(), nil, + grpcutil.ExpectedServiceDialOptions("resource_manager.ResourceManager")...) + re.NoError(err) + defer cc.Close() + stream, err := rmpb.NewResourceManagerClient(cc).AcquireTokenBuckets(ctx) + re.NoError(err) + // Wait for the stream to be rejected. + re.Eventually(func() bool { + return stream.Send(&rmpb.TokenBucketsRequest{}) != nil + }, 5*time.Second, 10*time.Millisecond) + + req := &tokenRequest{ + done: make(chan error, 1), + Request: &rmpb.TokenBucketsRequest{}, + } + err = (&innerClient{}).processTokenRequests(stream, req) + re.Equal(codes.Unimplemented, status.Code(errors.Unwrap(err))) + re.Equal(connectivity.Shutdown, cc.GetState()) + re.Equal(err, <-req.done) +} diff --git a/client/servicediscovery/expected_service.go b/client/servicediscovery/expected_service.go new file mode 100644 index 00000000000..b029a8ab98d --- /dev/null +++ b/client/servicediscovery/expected_service.go @@ -0,0 +1,37 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package servicediscovery + +import ( + "google.golang.org/grpc" + + "github.com/tikv/pd/client/pkg/utils/grpcutil" +) + +// Each discovery only checks the service which every endpoint it connects to +// must provide, so a connection to an address reused by another kind of +// service is closed and recreated instead of being reused forever. +var ( + pdExpectedServiceDialOptions = grpcutil.ExpectedServiceDialOptions("pdpb.PD") + tsoExpectedServiceDialOptions = grpcutil.ExpectedServiceDialOptions("tsopb.TSO") + routerExpectedServiceDialOptions = grpcutil.ExpectedServiceDialOptions("routerpb.Router") + resourceManagerExpectedServiceDialOptions = grpcutil.ExpectedServiceDialOptions("resource_manager.ResourceManager") +) + +func withExpectedService(dialOpts, expectedServiceOpts []grpc.DialOption) []grpc.DialOption { + opts := make([]grpc.DialOption, 0, len(dialOpts)+len(expectedServiceOpts)) + opts = append(opts, dialOpts...) + return append(opts, expectedServiceOpts...) +} diff --git a/client/servicediscovery/expected_service_test.go b/client/servicediscovery/expected_service_test.go new file mode 100644 index 00000000000..26d3a7fca46 --- /dev/null +++ b/client/servicediscovery/expected_service_test.go @@ -0,0 +1,282 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package servicediscovery + +import ( + "context" + "encoding/json" + "net" + "testing" + "time" + + "github.com/gogo/protobuf/proto" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/connectivity" + "google.golang.org/grpc/health" + healthpb "google.golang.org/grpc/health/grpc_health_v1" + "google.golang.org/grpc/status" + + "github.com/pingcap/kvproto/pkg/meta_storagepb" + "github.com/pingcap/kvproto/pkg/pdpb" + rmpb "github.com/pingcap/kvproto/pkg/resource_manager" + "github.com/pingcap/kvproto/pkg/routerpb" + "github.com/pingcap/kvproto/pkg/tsopb" + + "github.com/tikv/pd/client/clients/metastorage" + "github.com/tikv/pd/client/opt" +) + +// startOtherServiceServer starts a gRPC server which provides none of the PD +// client facing services, simulating an address reused by another service. +func startOtherServiceServer(t *testing.T) string { + lis, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + srv := grpc.NewServer() + healthpb.RegisterHealthServer(srv, health.NewServer()) + go func() { + _ = srv.Serve(lis) + }() + t.Cleanup(srv.Stop) + return "http://" + lis.Addr().String() +} + +func TestServiceDiscoveriesRecreateConnOnMissingService(t *testing.T) { + ctx := t.Context() + pdDiscovery := &serviceDiscovery{ctx: ctx, option: opt.NewOption()} + tsoDiscovery := &tsoServiceDiscovery{ctx: ctx, option: opt.NewOption()} + routerDiscovery := &routerServiceDiscovery{ctx: ctx, option: opt.NewOption()} + defer func() { + for _, d := range []ServiceDiscovery{pdDiscovery, tsoDiscovery, routerDiscovery} { + d.GetClientConns().Range(func(_, cc any) bool { + _ = cc.(*grpc.ClientConn).Close() + return true + }) + } + }() + + for _, tc := range []struct { + name string + discovery ServiceDiscovery + invoke func(context.Context, *grpc.ClientConn) error + }{ + { + name: "pd", + discovery: pdDiscovery, + invoke: func(ctx context.Context, cc *grpc.ClientConn) error { + _, err := pdpb.NewPDClient(cc).GetMembers(ctx, &pdpb.GetMembersRequest{}) + return err + }, + }, + { + name: "tso", + discovery: tsoDiscovery, + invoke: func(ctx context.Context, cc *grpc.ClientConn) error { + _, err := tsopb.NewTSOClient(cc).FindGroupByKeyspaceID(ctx, &tsopb.FindGroupByKeyspaceIDRequest{}) + return err + }, + }, + { + name: "router", + discovery: routerDiscovery, + invoke: func(ctx context.Context, cc *grpc.ClientConn) error { + _, err := routerpb.NewRouterClient(cc).GetRegion(ctx, &pdpb.GetRegionRequest{}) + return err + }, + }, + } { + t.Run(tc.name, func(t *testing.T) { + re := require.New(t) + url := startOtherServiceServer(t) + cc, err := tc.discovery.GetOrCreateGRPCConn(url) + re.NoError(err) + + ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + re.Equal(codes.Unimplemented, status.Code(tc.invoke(ctx, cc))) + re.Equal(connectivity.Shutdown, cc.GetState()) + + newCC, err := tc.discovery.GetOrCreateGRPCConn(url) + re.NoError(err) + re.NotSame(cc, newCC) + re.NotEqual(connectivity.Shutdown, newCC.GetState()) + }) + } +} + +func TestResourceManagerDiscoveryRecreatesConnOnMissingService(t *testing.T) { + re := require.New(t) + ctx := t.Context() + url := startOtherServiceServer(t) + + var notifications int + discovery := NewResourceManagerDiscovery(ctx, 1, nil, nil, opt.NewOption(), func(string) error { + notifications++ + return nil + }) + defer discovery.Close() + discovery.resetConn(url) + re.Equal(1, notifications) + cc := discovery.GetConn() + re.NotNil(cc) + + callCtx, callCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer callCancel() + _, err := rmpb.NewResourceManagerClient(cc).ListResourceGroups(callCtx, &rmpb.ListResourceGroupsRequest{}) + re.Equal(codes.Unimplemented, status.Code(err)) + re.Equal(connectivity.Shutdown, cc.GetState()) + + // The service URL is unchanged, but the closed connection must be recreated. + discovery.resetConn(url) + newCC := discovery.GetConn() + re.NotNil(newCC) + re.NotSame(cc, newCC) + re.NotEqual(connectivity.Shutdown, newCC.GetState()) + re.Equal(url, discovery.GetServiceURL()) + // The token stream on the closed connection must switch to the new one. + re.Equal(2, notifications) +} + +// fixedRevisionMetaStorageClient always returns the same revision, i.e., the +// resource manager primary is unchanged. +type fixedRevisionMetaStorageClient struct { + countingMetaStorageClient +} + +func (c *fixedRevisionMetaStorageClient) Get(context.Context, []byte, ...opt.MetaStorageOption) (*meta_storagepb.GetResponse, error) { + return &meta_storagepb.GetResponse{ + Header: &meta_storagepb.ResponseHeader{Revision: 1}, + Kvs: []*meta_storagepb.KeyValue{{Value: c.value}}, + Count: 1, + }, nil +} + +func TestResourceManagerDiscoveryUpdateRecreatesClosedConn(t *testing.T) { + re := require.New(t) + oldRetryInterval := serviceURLRetryInterval + serviceURLRetryInterval = 10 * time.Millisecond + defer func() { + serviceURLRetryInterval = oldRetryInterval + }() + url := startOtherServiceServer(t) + value, err := proto.Marshal(&rmpb.Participant{ListenUrls: []string{url}}) + re.NoError(err) + metaCli := &fixedRevisionMetaStorageClient{countingMetaStorageClient{value: value}} + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + discovery := NewResourceManagerDiscovery(ctx, 1, metaCli, nil, opt.NewOption(), func(string) error { return nil }) + discovery.Init() + defer discovery.Close() + re.Eventually(func() bool { + return discovery.GetConn() != nil + }, 5*time.Second, 10*time.Millisecond) + cc := discovery.GetConn() + re.NoError(cc.Close()) + + discovery.ScheduleUpdateServiceURL() + re.Eventually(func() bool { + newCC := discovery.GetConn() + return newCC != nil && newCC != cc && newCC.GetState() != connectivity.Shutdown + }, 5*time.Second, 10*time.Millisecond) +} + +// staticRegistryMetaStorageClient returns the same router service registry. +type staticRegistryMetaStorageClient struct { + countingMetaStorageClient +} + +func (c *staticRegistryMetaStorageClient) Get(context.Context, []byte, ...opt.MetaStorageOption) (*meta_storagepb.GetResponse, error) { + return &meta_storagepb.GetResponse{ + Header: &meta_storagepb.ResponseHeader{Revision: 1}, + Kvs: []*meta_storagepb.KeyValue{{Value: c.value}}, + Count: 1, + }, nil +} + +func TestRouterServiceDiscoveryRecreatesClosedNodeConn(t *testing.T) { + re := require.New(t) + url := startOtherServiceServer(t) + value, err := json.Marshal(&metastorage.ServiceRegistryEntry{ServiceAddr: url}) + re.NoError(err) + + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + discovery := &routerServiceDiscovery{ + ServiceDiscovery: NewMockServiceDiscovery(nil, nil), + ctx: ctx, + cancel: cancel, + metaCli: &staticRegistryMetaStorageClient{countingMetaStorageClient{value: value}}, + option: opt.NewOption(), + checkMembershipCh: make(chan struct{}, 1), + balancer: newServiceBalancer(emptyErrorFn), + callbacks: newServiceCallbacks(), + } + defer discovery.Close() + re.NoError(discovery.updateMember()) + node, ok := discovery.nodes.Load(url) + re.True(ok) + cc := node.(*serviceClient).GetClientConn() + re.NotNil(cc) + re.NoError(cc.Close()) + + // The router service URLs are unchanged, but the closed connection must be recreated. + re.NoError(discovery.updateMember()) + node, ok = discovery.nodes.Load(url) + re.True(ok) + newCC := node.(*serviceClient).GetClientConn() + re.NotNil(newCC) + re.NotSame(cc, newCC) + conns := 0 + discovery.GetClientConns().Range(func(_, value any) bool { + conns++ + re.Same(newCC, value) + return true + }) + re.Equal(1, conns) +} + +func TestServiceDiscoverySwitchLeaderRecreatesClosedConn(t *testing.T) { + re := require.New(t) + url := startOtherServiceServer(t) + ctx := t.Context() + discovery := &serviceDiscovery{ctx: ctx, option: opt.NewOption(), callbacks: newServiceCallbacks()} + defer discovery.clientConns.Range(func(_, cc any) bool { + _ = cc.(*grpc.ClientConn).Close() + return true + }) + discovery.leader.Store(&serviceClient{}) + + changed, err := discovery.switchLeader(url) + re.NoError(err) + re.True(changed) + cc := discovery.GetServiceClient().GetClientConn() + re.NotNil(cc) + + callCtx, callCancel := context.WithTimeout(context.Background(), 5*time.Second) + defer callCancel() + _, err = pdpb.NewPDClient(cc).GetMembers(callCtx, &pdpb.GetMembersRequest{}) + re.Equal(codes.Unimplemented, status.Code(err)) + re.Nil(discovery.GetServiceClient().GetClientConn()) + + // The leader is unchanged, but the closed connection must be recreated. + changed, err = discovery.switchLeader(url) + re.NoError(err) + re.True(changed) + newCC := discovery.GetServiceClient().GetClientConn() + re.NotNil(newCC) + re.NotSame(cc, newCC) +} diff --git a/client/servicediscovery/resource_manager_service_discovery.go b/client/servicediscovery/resource_manager_service_discovery.go index 53b521f3d3e..9e4f2fc649d 100644 --- a/client/servicediscovery/resource_manager_service_discovery.go +++ b/client/servicediscovery/resource_manager_service_discovery.go @@ -25,6 +25,7 @@ import ( "github.com/gogo/protobuf/proto" "go.uber.org/zap" "google.golang.org/grpc" + "google.golang.org/grpc/connectivity" "github.com/pingcap/kvproto/pkg/resource_manager" "github.com/pingcap/log" @@ -121,7 +122,9 @@ func (r *ResourceManagerDiscovery) resetConn(url string) { r.mu.Lock() defer r.mu.Unlock() // An unchanged empty URL still means PD fallback, not a leader change. - if r.serviceURL == url { + // An unchanged URL with a closed connection, e.g., the address was reused by + // another service, still needs to reconnect. + if r.serviceURL == url && (len(url) == 0 || !isConnShutdown(r.conn)) { return } if len(url) == 0 { @@ -133,7 +136,8 @@ func (r *ResourceManagerDiscovery) resetConn(url string) { _ = r.onLeaderChanged("") return } - newConn, err := grpcutil.GetClientConn(r.ctx, url, r.tlsCfg, r.option.GRPCDialOptions...) + newConn, err := grpcutil.GetClientConn(r.ctx, url, r.tlsCfg, + withExpectedService(r.option.GRPCDialOptions, resourceManagerExpectedServiceDialOptions)...) if err != nil { // Dial without `WithBlock`, normally it should not fail. log.Error("[resource-manager] failed to create gRPC connection", @@ -166,6 +170,16 @@ func (r *ResourceManagerDiscovery) GetConn() *grpc.ClientConn { return r.conn } +func (r *ResourceManagerDiscovery) isConnShutdown() bool { + r.mu.RLock() + defer r.mu.RUnlock() + return len(r.serviceURL) != 0 && isConnShutdown(r.conn) +} + +func isConnShutdown(conn *grpc.ClientConn) bool { + return conn != nil && conn.GetState() == connectivity.Shutdown +} + func (r *ResourceManagerDiscovery) discoverServiceURL() (string, int64, error) { resp, err := r.metaCli.Get(r.ctx, []byte(r.discoveryKey)) if err != nil { @@ -230,9 +244,9 @@ func (r *ResourceManagerDiscovery) updateServiceURLLoop(revision int64) { zap.Error(err)) return } - if newRevision > revision { + if newRevision > revision || r.isConnShutdown() { r.resetConn(url) - revision = newRevision + revision = max(revision, newRevision) } } for { diff --git a/client/servicediscovery/router_service_discovery.go b/client/servicediscovery/router_service_discovery.go index eabcb6db687..27024700f59 100644 --- a/client/servicediscovery/router_service_discovery.go +++ b/client/servicediscovery/router_service_discovery.go @@ -101,7 +101,8 @@ func (r *routerServiceDiscovery) GetServiceClientByKind(_ APIKind) ServiceClient // GetOrCreateGRPCConn creates a gRPC connection to the router service. func (r *routerServiceDiscovery) GetOrCreateGRPCConn(url string) (*grpc.ClientConn, error) { - return grpcutil.GetOrCreateGRPCConn(r.ctx, &r.clientConns, url, r.tlsCfg, r.option.GRPCDialOptions...) + return grpcutil.GetOrCreateGRPCConn(r.ctx, &r.clientConns, url, r.tlsCfg, + withExpectedService(r.option.GRPCDialOptions, routerExpectedServiceDialOptions)...) } // RemoveClientConn removes and closes the grpc client connection of the given URL. @@ -196,14 +197,26 @@ func (r *routerServiceDiscovery) updateMember() error { return err } changed := r.nodesChanged(sortedAddrs) - if !changed { + if changed { + r.updateURLs(sortedAddrs) + } else if !r.hasClosedNodeConn() { return nil } - r.updateURLs(sortedAddrs) + // Recreate the closed connections even if the nodes are unchanged, e.g., the + // address was reused by another service. r.updateNodes(sortedAddrs) return nil } +func (r *routerServiceDiscovery) hasClosedNodeConn() bool { + closed := false + r.nodes.Range(func(_, value any) bool { + closed = value.(*serviceClient).GetClientConn() == nil + return !closed + }) + return closed +} + func (r *routerServiceDiscovery) updateNodes(urls []string) { urlSet := make(map[string]struct{}, len(urls)) for _, url := range urls { @@ -222,7 +235,7 @@ func (r *routerServiceDiscovery) updateNodes(urls []string) { continue } nodeClient := newPDServiceClient(newURL, r.GetServingURL(), conn, false) - r.nodes.LoadOrStore(newURL, nodeClient) + r.nodes.Store(newURL, nodeClient) } } clients := make([]ServiceClient, 0, len(urls)) diff --git a/client/servicediscovery/service_discovery.go b/client/servicediscovery/service_discovery.go index 05e6a638f77..58585dbf680 100644 --- a/client/servicediscovery/service_discovery.go +++ b/client/servicediscovery/service_discovery.go @@ -1110,7 +1110,8 @@ func (c *serviceDiscovery) updateServiceClient(members []*pdpb.Member, leader *p // GetOrCreateGRPCConn returns the corresponding grpc client connection of the given URL. func (c *serviceDiscovery) GetOrCreateGRPCConn(url string) (*grpc.ClientConn, error) { - return grpcutil.GetOrCreateGRPCConn(c.ctx, &c.clientConns, url, c.tlsCfg, c.option.GRPCDialOptions...) + return grpcutil.GetOrCreateGRPCConn(c.ctx, &c.clientConns, url, c.tlsCfg, + withExpectedService(c.option.GRPCDialOptions, pdExpectedServiceDialOptions)...) } // RemoveClientConn removes and closes the grpc client connection of the given URL. diff --git a/client/servicediscovery/tso_service_discovery.go b/client/servicediscovery/tso_service_discovery.go index aa732cf1d2f..9505368c920 100644 --- a/client/servicediscovery/tso_service_discovery.go +++ b/client/servicediscovery/tso_service_discovery.go @@ -327,7 +327,8 @@ func (c *tsoServiceDiscovery) GetBackupURLs() []string { // GetOrCreateGRPCConn returns the corresponding grpc client connection of the given URL. func (c *tsoServiceDiscovery) GetOrCreateGRPCConn(url string) (*grpc.ClientConn, error) { - return grpcutil.GetOrCreateGRPCConn(c.ctx, &c.clientConns, url, c.tlsCfg, c.option.GRPCDialOptions...) + return grpcutil.GetOrCreateGRPCConn(c.ctx, &c.clientConns, url, c.tlsCfg, + withExpectedService(c.option.GRPCDialOptions, tsoExpectedServiceDialOptions)...) } // RemoveClientConn removes and closes the gRPC connection of the given URL. diff --git a/tests/integrations/client/stale_address_test.go b/tests/integrations/client/stale_address_test.go new file mode 100644 index 00000000000..133ee4cdf22 --- /dev/null +++ b/tests/integrations/client/stale_address_test.go @@ -0,0 +1,128 @@ +// Copyright 2026 TiKV Project Authors. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package client_test + +import ( + "context" + "net" + "strings" + "sync" + "sync/atomic" + "testing" + + "github.com/stretchr/testify/require" + "google.golang.org/grpc" + "google.golang.org/grpc/health" + healthpb "google.golang.org/grpc/health/grpc_health_v1" + + "github.com/tikv/pd/client/opt" + "github.com/tikv/pd/pkg/utils/testutil" + "github.com/tikv/pd/tests" +) + +// staleAddressDialer simulates a stale DNS record: once redirected, dialing the +// PD address reaches another service which reused the address. +type staleAddressDialer struct { + redirectAddr atomic.Pointer[string] + + mu sync.Mutex + conns []net.Conn +} + +func (d *staleAddressDialer) dial(ctx context.Context, addr string) (net.Conn, error) { + if redirectAddr := d.redirectAddr.Load(); redirectAddr != nil { + addr = *redirectAddr + } + conn, err := (&net.Dialer{}).DialContext(ctx, "tcp", addr) + if err != nil { + return nil, err + } + d.mu.Lock() + defer d.mu.Unlock() + d.conns = append(d.conns, conn) + return conn, nil +} + +// redirect redirects the following dials to the given address, and breaks the +// established connections to simulate the old PD instance going away. +func (d *staleAddressDialer) redirect(addr string) { + d.redirectAddr.Store(&addr) + d.mu.Lock() + defer d.mu.Unlock() + for _, conn := range d.conns { + _ = conn.Close() + } + d.conns = nil +} + +// recover makes the following dials reach the PD address again, i.e., the DNS +// record is refreshed. +func (d *staleAddressDialer) recover() { + d.redirectAddr.Store(nil) +} + +// startNonPDServer starts a gRPC server which does not provide the PD service, +// e.g., a scheduling service instance which reused the address of a PD. +func startNonPDServer(t *testing.T) string { + lis, err := net.Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + srv := grpc.NewServer() + healthpb.RegisterHealthServer(srv, health.NewServer()) + go func() { + _ = srv.Serve(lis) + }() + t.Cleanup(srv.Stop) + return lis.Addr().String() +} + +func TestClientRecoversFromAddressReusedByOtherService(t *testing.T) { + re := require.New(t) + ctx := t.Context() + cluster, err := tests.NewTestCluster(ctx, 1) + re.NoError(err) + defer cluster.Destroy() + endpoints := runServer(re, cluster) + + dialer := &staleAddressDialer{} + cli := setupCli(ctx, re, endpoints, + opt.WithGRPCDialOptions(grpc.WithContextDialer(dialer.dial))) + defer cli.Close() + + requestsSucceed := func() bool { + if _, _, err := cli.GetTS(ctx); err != nil { + t.Log(err) + return false + } + if _, err := cli.GetAllStores(ctx); err != nil { + t.Log(err) + return false + } + return true + } + testutil.Eventually(re, requestsSucceed) + + // The PD address is reused by another service, and the stale DNS record + // makes the client reconnect to it. + dialer.redirect(startNonPDServer(t)) + testutil.Eventually(re, func() bool { + _, err := cli.GetAllStores(ctx) + return err != nil && strings.Contains(err.Error(), "unknown service pdpb.PD") + }) + + // After the DNS record is refreshed, the client must not keep using the + // connection to the wrong service. + dialer.recover() + testutil.Eventually(re, requestsSucceed) +}