pkg: support service discovery

pkg/mcs/discovery/discover.go

@@ -0,0 +1,37 @@
+// Copyright 2023 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 discovery
+
+import (
+	"github.com/tikv/pd/pkg/utils/etcdutil"
+	"go.etcd.io/etcd/clientv3"
+)
+
+// Discover is used to get all the service instances of the specified service name.
+func Discover(cli *clientv3.Client, serviceName string) ([]string, error) {
+	key := discoveryPath(serviceName) + "/"
+	endKey := clientv3.GetPrefixRangeEnd(key) + "/"
+
+	withRange := clientv3.WithRange(endKey)
+	resp, err := etcdutil.EtcdKVGet(cli, key, withRange)
+	if err != nil {
+		return nil, err
+	}
+	values := make([]string, 0, len(resp.Kvs))
+	for _, item := range resp.Kvs {
+		values = append(values, string(item.Value))
+	}
+	return values, nil
+}

pkg/mcs/discovery/discover_test.go

@@ -0,0 +1,63 @@
+// Copyright 2023 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 discovery
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/utils/etcdutil"
+	"go.etcd.io/etcd/clientv3"
+	"go.etcd.io/etcd/embed"
+)
+
+func TestDiscover(t *testing.T) {
+	re := require.New(t)
+	cfg := etcdutil.NewTestSingleConfig(t)
+	etcd, err := embed.StartEtcd(cfg)
+	defer func() {
+		etcd.Close()
+	}()
+	re.NoError(err)
+
+	ep := cfg.LCUrls[0].String()
+	client, err := clientv3.New(clientv3.Config{
+		Endpoints: []string{ep},
+	})
+	re.NoError(err)
+
+	<-etcd.Server.ReadyNotify()
+	sr1 := NewServiceRegister(context.Background(), client, "test_service", "127.0.0.1:1", "127.0.0.1:1", 1)
+	err = sr1.Register()
+	re.NoError(err)
+	sr2 := NewServiceRegister(context.Background(), client, "test_service", "127.0.0.1:2", "127.0.0.1:2", 1)
+	err = sr2.Register()
+	re.NoError(err)
+
+	endpoints, err := Discover(client, "test_service")
+	re.NoError(err)
+	re.Len(endpoints, 2)
+	re.Equal("127.0.0.1:1", endpoints[0])
+	re.Equal("127.0.0.1:2", endpoints[1])
+
+	sr1.cancel()
+	sr2.cancel()
+	time.Sleep(3 * time.Second)
+	endpoints, err = Discover(client, "test_service")
+	re.NoError(err)
+	re.Empty(endpoints)
+}

pkg/mcs/discovery/key_path.go

@@ -0,0 +1,30 @@
+// Copyright 2023 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 discovery
+
+import "path"
+
+const (
+	registryPrefix = "/pd/microservice"
+	registryKey    = "registry"
+)
+
+func registryPath(serviceName, serviceAddr string) string {
+	return path.Join(registryPrefix, serviceName, registryKey, serviceAddr)
+}
+
+func discoveryPath(serviceName string) string {
+	return path.Join(registryPrefix, serviceName, registryKey)
+}

pkg/mcs/discovery/register.go

@@ -0,0 +1,108 @@
+// Copyright 2023 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 discovery
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"github.com/pingcap/log"
+	"go.etcd.io/etcd/clientv3"
+	"go.uber.org/zap"
+)
+
+// ServiceRegister is used to register the service to etcd.
+type ServiceRegister struct {
+	ctx    context.Context
+	cancel context.CancelFunc
+	cli    *clientv3.Client
+	key    string
+	value  string
+	ttl    int64
+}
+
+// NewServiceRegister creates a new ServiceRegister.
+func NewServiceRegister(ctx context.Context, cli *clientv3.Client, serviceName, serviceAddr, serializedValue string, ttl int64) *ServiceRegister {
+	cctx, cancel := context.WithCancel(ctx)
+	serviceKey := registryPath(serviceName, serviceAddr)
+	return &ServiceRegister{
+		ctx:    cctx,
+		cancel: cancel,
+		cli:    cli,
+		key:    serviceKey,
+		value:  serializedValue,
+		ttl:    ttl,
+	}
+}
+
+// Register registers the service to etcd.
+func (sr *ServiceRegister) Register() error {
+	resp, err := sr.cli.Grant(sr.ctx, sr.ttl)
+	if err != nil {
+		sr.cancel()
+		return fmt.Errorf("grant lease failed: %v", err)
+	}
+
+	if _, err := sr.cli.Put(sr.ctx, sr.key, sr.value, clientv3.WithLease(resp.ID)); err != nil {
+		sr.cancel()
+		return fmt.Errorf("put the key %s failed: %v", sr.key, err)
+	}
+
+	kresp, err := sr.cli.KeepAlive(sr.ctx, resp.ID)
+	if err != nil {
+		sr.cancel()
+		return fmt.Errorf("keepalive failed: %v", err)
+	}
+	go func() {
+		for {
+			select {
+			case <-sr.ctx.Done():
+				log.Info("exit register process", zap.String("key", sr.key))
+				return
+			case _, ok := <-kresp:
+				if !ok {
+					log.Error("keep alive failed", zap.String("key", sr.key))
+					// retry
+					t := time.NewTicker(time.Duration(sr.ttl) * time.Second / 2)
+					for {
+						<-t.C
+						resp, err := sr.cli.Grant(sr.ctx, sr.ttl)
+						if err != nil {
+							log.Error("grant lease failed", zap.String("key", sr.key), zap.Error(err))
+							continue
+						}
+
+						if _, err := sr.cli.Put(sr.ctx, sr.key, sr.value, clientv3.WithLease(resp.ID)); err != nil {
+							log.Error("put the key failed", zap.String("key", sr.key), zap.Error(err))
+							continue
+						}
+					}
+				}
+			}
+		}
+	}()
+
+	return nil
+}
+
+// Deregister deregisters the service from etcd.
+func (sr *ServiceRegister) Deregister() error {
+	sr.cancel()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	_, err := sr.cli.Delete(ctx, sr.key)
+	return err
+}

pkg/mcs/discovery/register_test.go

@@ -0,0 +1,66 @@
+// Copyright 2023 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 discovery
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/require"
+	"github.com/tikv/pd/pkg/utils/etcdutil"
+	"go.etcd.io/etcd/clientv3"
+	"go.etcd.io/etcd/embed"
+)
+
+func TestRegister(t *testing.T) {
+	re := require.New(t)
+	cfg := etcdutil.NewTestSingleConfig(t)
+	etcd, err := embed.StartEtcd(cfg)
+	defer func() {
+		etcd.Close()
+	}()
+	re.NoError(err)
+
+	ep := cfg.LCUrls[0].String()
+	client, err := clientv3.New(clientv3.Config{
+		Endpoints: []string{ep},
+	})
+	re.NoError(err)
+
+	<-etcd.Server.ReadyNotify()
+	sr := NewServiceRegister(context.Background(), client, "test_service", "127.0.0.1:1", "127.0.0.1:1", 10)
+	err = sr.Register()
+	re.NoError(err)
+	resp, err := client.Get(context.Background(), sr.key)
+	re.NoError(err)
+	re.Equal("127.0.0.1:1", string(resp.Kvs[0].Value))
+
+	err = sr.Deregister()
+	re.NoError(err)
+	resp, err = client.Get(context.Background(), sr.key)
+	re.NoError(err)
+	re.Empty(resp.Kvs)
+
+	sr = NewServiceRegister(context.Background(), client, "test_service", "127.0.0.1:2", "127.0.0.1:2", 1)
+	err = sr.Register()
+	re.NoError(err)
+	sr.cancel()
+	// ensure that the lease is expired
+	time.Sleep(3 * time.Second)
+	resp, err = client.Get(context.Background(), sr.key)
+	re.NoError(err)
+	re.Empty(resp.Kvs)
+}

server/schedule/operator/builder.go

@@ -523,8 +523,8 @@ func (b *Builder) prepareBuild() (string, error) {
 		}
 	}
 
-	// Although switch witness may have nothing to do with conf change (except switch witness voter to non-witness voter:
-	// it will domote to learner first, then switch witness, finally promote the non-witness learner to voter back),
+	// Although switching witness may have nothing to do with conf change (except switch witness voter to non-witness voter:
+	// it will demote to learner first, then switch witness, finally promote the non-witness learner to voter back),
 	// the logic here is reused for batch switch.
 	if len(b.toAdd)+len(b.toRemove)+len(b.toPromote)+len(b.toWitness)+len(b.toNonWitness)+len(b.toPromoteNonWitness) <= 1 &&
 		len(b.toDemote) == 0 && !(len(b.toRemove) == 1 && len(b.targetPeers) == 1) {