From cbfbb3a4f86fcdb3bfab4b9a148184072d9fb0f3 Mon Sep 17 00:00:00 2001
From: Ryan Leung <rleungx@gmail.com>
Date: Fri, 8 Mar 2024 17:21:16 +0800
Subject: [PATCH] prevent keyspace client from panic

Signed-off-by: Ryan Leung <rleungx@gmail.com>
---
 client/keyspace_client.go | 22 +++++++++++++++++++---
 1 file changed, 19 insertions(+), 3 deletions(-)

diff --git a/client/keyspace_client.go b/client/keyspace_client.go
index 80b9cb0e436..340ecd0250e 100644
--- a/client/keyspace_client.go
+++ b/client/keyspace_client.go
@@ -21,6 +21,7 @@ import (
 	"github.com/opentracing/opentracing-go"
 	"github.com/pingcap/errors"
 	"github.com/pingcap/kvproto/pkg/keyspacepb"
+	"github.com/tikv/pd/client/errs"
 )
 
 // KeyspaceClient manages keyspace metadata.
@@ -56,7 +57,12 @@ func (c *client) LoadKeyspace(ctx context.Context, name string) (*keyspacepb.Key
 		Header: c.requestHeader(),
 		Name:   name,
 	}
-	resp, err := c.keyspaceClient().LoadKeyspace(ctx, req)
+	protoClient := c.keyspaceClient()
+	if protoClient == nil {
+		cancel()
+		return nil, errs.ErrClientGetProtoClient
+	}
+	resp, err := protoClient.LoadKeyspace(ctx, req)
 	cancel()
 
 	if err != nil {
@@ -96,7 +102,12 @@ func (c *client) UpdateKeyspaceState(ctx context.Context, id uint32, state keysp
 		Id:     id,
 		State:  state,
 	}
-	resp, err := c.keyspaceClient().UpdateKeyspaceState(ctx, req)
+	protoClient := c.keyspaceClient()
+	if protoClient == nil {
+		cancel()
+		return nil, errs.ErrClientGetProtoClient
+	}
+	resp, err := protoClient.UpdateKeyspaceState(ctx, req)
 	cancel()
 
 	if err != nil {
@@ -135,7 +146,12 @@ func (c *client) GetAllKeyspaces(ctx context.Context, startID uint32, limit uint
 		StartId: startID,
 		Limit:   limit,
 	}
-	resp, err := c.keyspaceClient().GetAllKeyspaces(ctx, req)
+	protoClient := c.keyspaceClient()
+	if protoClient == nil {
+		cancel()
+		return nil, errs.ErrClientGetProtoClient
+	}
+	resp, err := protoClient.GetAllKeyspaces(ctx, req)
 	cancel()
 
 	if err != nil {