Repository: pingcap/tidb
File: pkg/meta/autoid/autoid_service.go (Go)

Code before:
func (d *ClientDiscover) GetClient(ctx context.Context, keyspaceID uint32) (autoid.AutoIDAllocClient, uint64, error) {
	d.mu.RLock()
	cli := d.mu.AutoIDAllocClient
	if cli != nil {
		d.mu.RUnlock()
		return cli, atomic.LoadUint64(&d.version), nil
	}
	d.mu.RUnlock()

	d.mu.Lock()
	defer d.mu.Unlock()
	if d.mu.AutoIDAllocClient != nil {
		return d.mu.AutoIDAllocClient, atomic.LoadUint64(&d.version), nil
	}
	// write a for loop to retry in case of etcd connection error.
	var resp *clientv3.GetResponse
	var err error
	for range []int{0, 1} {
		resp, err = d.etcdCli.Get(ctx, GetAutoIDServiceLeaderEtcdPath(keyspaceID), clientv3.WithFirstCreate()...)
		if err != nil {
			return nil, 0, errors.Trace(err)
		}
		if len(resp.Kvs) == 0 {
			// If the key is not found, it means the autoid service leader is not elected yet.
			// We can retry to get the leader.
			time.Sleep(10 * time.Millisecond)
			continue
		}
		break
	}
	if len(resp.Kvs) == 0 {
		return nil, 0, errors.New("autoid service leader not found")
	}

	addr := string(resp.Kvs[0].Value)
	opt := grpc.WithTransportCredentials(insecure.NewCredentials())
	security := config.GetGlobalConfig().Security
	if len(security.ClusterSSLCA) != 0 {
		clusterSecurity := security.ClusterSecurity()
		tlsConfig, err := clusterSecurity.ToTLSConfig()
		if err != nil {
			return nil, 0, errors.Trace(err)
		}
		opt = grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))
	}
	logutil.BgLogger().Info("connect to leader", zap.String("category", "autoid client"), zap.String("addr", addr))
	grpcConn, err := grpc.NewClient(addr, opt)
	if err != nil {
		return nil, 0, errors.Trace(err)
	}
	cli = autoid.NewAutoIDAllocClient(grpcConn)

Review comment (suggestion):
please use `select { <- ctx.Done(): ... <- time.After(...): ... } ` to find context canceled earlier.

Code after:
func (d *ClientDiscover) GetClient(ctx context.Context, keyspaceID uint32) (autoid.AutoIDAllocClient, uint64, error) {
	d.mu.RLock()
	cli := d.mu.AutoIDAllocClient
	if cli != nil {
		d.mu.RUnlock()
		return cli, atomic.LoadUint64(&d.version), nil
	}
	d.mu.RUnlock()

	d.mu.Lock()
	defer d.mu.Unlock()
	if d.mu.AutoIDAllocClient != nil {
		return d.mu.AutoIDAllocClient, atomic.LoadUint64(&d.version), nil
	}
	// write a for loop to retry in case of etcd connection error.
	var resp *clientv3.GetResponse
	var err error
	var bo backoffer
retry:
	resp, err = d.etcdCli.Get(ctx, GetAutoIDServiceLeaderEtcdPath(keyspaceID), clientv3.WithFirstCreate()...)
	if err != nil {
		return nil, 0, errors.Trace(err)
	}
	if len(resp.Kvs) == 0 {
		// If the key is not found, it means the autoid service leader is not elected yet.
		// We can retry to get the leader.
		if err := ctx.Err(); err != nil {
			return nil, 0, errors.Trace(err)
		}
		bo.Backoff()
		goto retry
	}
	bo.Reset()

	addr := string(resp.Kvs[0].Value)
	opt := grpc.WithTransportCredentials(insecure.NewCredentials())
	security := config.GetGlobalConfig().Security
	if len(security.ClusterSSLCA) != 0 {
		clusterSecurity := security.ClusterSecurity()
		tlsConfig, err := clusterSecurity.ToTLSConfig()
		if err != nil {
			return nil, 0, errors.Trace(err)
		}
		opt = grpc.WithTransportCredentials(credentials.NewTLS(tlsConfig))
	}
	logutil.BgLogger().Info("connect to leader", zap.String("category", "autoid client"), zap.String("addr", addr))
	grpcConn, err := grpc.NewClient(addr, opt)
	if err != nil {
		return nil, 0, errors.Trace(err)
	}
	cli = autoid.NewAutoIDAllocClient(grpcConn)