From f2e018036ad2bdaf45b66ebcc5abe854650e84a8 Mon Sep 17 00:00:00 2001
From: Ryan Leung <rleungx@gmail.com>
Date: Fri, 5 May 2023 15:57:56 +0800
Subject: [PATCH] client: retry load keyspace meta when creating a new client
 (#6402)

ref tikv/pd#5895

Signed-off-by: Ryan Leung <rleungx@gmail.com>

Co-authored-by: ti-chi-bot[bot] <108142056+ti-chi-bot[bot]@users.noreply.github.com>
---
 client/client.go | 26 ++++++++++++++++++++++++--
 1 file changed, 24 insertions(+), 2 deletions(-)

diff --git a/client/client.go b/client/client.go
index b3885c53791..d9cb0358f3f 100644
--- a/client/client.go
+++ b/client/client.go
@@ -383,13 +383,35 @@ func NewClientWithKeyspaceName(ctx context.Context, keyspace string, svrAddrs []
 		c.cancel()
 		return nil, err
 	}
+	if err := c.initRetry(c.loadKeyspaceMeta, keyspace); err != nil {
+		return nil, err
+	}
+	return c, nil
+}
+
+func (c *client) initRetry(f func(s string) error, str string) error {
+	var err error
+	for i := 0; i < c.option.maxRetryTimes; i++ {
+		if err = f(str); err == nil || strings.Contains(err.Error(), "ENTRY_NOT_FOUND") {
+			return nil
+		}
+		select {
+		case <-c.ctx.Done():
+			return err
+		case <-time.After(time.Second):
+		}
+	}
+	return errors.WithStack(err)
+}
+
+func (c *client) loadKeyspaceMeta(keyspace string) error {
 	keyspaceMeta, err := c.LoadKeyspace(context.TODO(), keyspace)
 	// Here we ignore ENTRY_NOT_FOUND error and it will set the keyspaceID to 0.
 	if err != nil && !strings.Contains(err.Error(), "ENTRY_NOT_FOUND") {
-		return nil, err
+		return err
 	}
 	c.keyspaceID = keyspaceMeta.GetId()
-	return c, nil
+	return nil
 }
 
 func (c *client) setup() error {