ha: refactor the schedule model

dm/common/common.go

@@ -25,16 +25,16 @@ import (
 var (
 	useOfClosedErrMsg = "use of closed network connection"
 	// WorkerRegisterKeyAdapter used to encode and decode register key.
-	// k/v: Encode(addr) -> name
+	// k/v: Encode(name) -> the information of the DM-worker node.
 	WorkerRegisterKeyAdapter KeyAdapter = keyHexEncoderDecoder("/dm-worker/r/")
 	// WorkerKeepAliveKeyAdapter used to encode and decode keepalive key.
-	// k/v: Encode(addr,name) -> time
+	// k/v: Encode(name) -> time
 	WorkerKeepAliveKeyAdapter KeyAdapter = keyHexEncoderDecoder("/dm-worker/a/")
 	// UpstreamConfigKeyAdapter store all config of which MySQL-task has not stopped.
 	// k/v: Encode(source-id) -> config
 	UpstreamConfigKeyAdapter KeyAdapter = keyEncoderDecoder("/dm-master/upstream/config/")
 	// UpstreamBoundWorkerKeyAdapter used to store address of worker in which MySQL-tasks which are running.
-	// k/v: Encode(addr) -> source-id
+	// k/v: Encode(name) -> the bound relationship.
 	UpstreamBoundWorkerKeyAdapter KeyAdapter = keyHexEncoderDecoder("/dm-master/bound-worker/")
 	// UpstreamSubTaskKeyAdapter used to store SubTask which are subscribing data from MySQL source.
 	// k/v: Encode(source-id, task-name) -> SubTaskConfig
@@ -50,9 +50,9 @@ var (
 
 func keyAdapterKeysLen(s KeyAdapter) int {
 	switch s {
-	case WorkerRegisterKeyAdapter, UpstreamConfigKeyAdapter, UpstreamBoundWorkerKeyAdapter:
+	case WorkerRegisterKeyAdapter, UpstreamConfigKeyAdapter, UpstreamBoundWorkerKeyAdapter, WorkerKeepAliveKeyAdapter:
 		return 1
-	case WorkerKeepAliveKeyAdapter, UpstreamSubTaskKeyAdapter:
+	case UpstreamSubTaskKeyAdapter:
 		return 2
 	}
 	return -1

dm/common/common_test.go

@@ -39,9 +39,9 @@ func (t *testCommon) TestKeyAdapter(c *C) {
 			want:    "/dm-worker/r/3132372e302e302e313a32333832",
 		},
 		{
-			keys:    []string{"127.0.0.1:2382", "worker1"},
+			keys:    []string{"worker1"},
 			adapter: WorkerKeepAliveKeyAdapter,
-			want:    "/dm-worker/a/3132372e302e302e313a32333832/776f726b657231",
+			want:    "/dm-worker/a/776f726b657231",
 		},
 		{
 			keys:    []string{"mysql1"},

pkg/ha/bound.go

@@ -0,0 +1,176 @@
+// Copyright 2020 PingCAP, Inc.
+//
+// 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,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ha
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+
+	"go.etcd.io/etcd/clientv3"
+	"go.etcd.io/etcd/mvcc/mvccpb"
+	"go.uber.org/zap"
+
+	"github.com/pingcap/dm/dm/common"
+	"github.com/pingcap/dm/pkg/etcdutil"
+	"github.com/pingcap/dm/pkg/log"
+)
+
+// SourceBound represents the bound relationship between the DM-worker instance and the upstream MySQL source.
+type SourceBound struct {
+	Source string `json:"source"` // the source ID of the upstream.
+	Worker string `json:"worker"` // the name of the bounded DM-worker for the source.
+}
+
+// NewSourceBound creates a new SourceBound instance.
+func NewSourceBound(source, worker string) SourceBound {
+	return SourceBound{
+		Source: source,
+		Worker: worker,
+	}
+}
+
+// NotBound returns whether the relationship has not bound.
+// An empty bound means the relationship has not bound.
+func (b SourceBound) NotBound() bool {
+	return b.Source == "" && b.Worker == ""
+}
+
+// String implements Stringer interface.
+func (b SourceBound) String() string {
+	s, _ := b.toJSON()
+	return s
+}
+
+// toJSON returns the string of JSON represent.
+func (b SourceBound) toJSON() (string, error) {
+	data, err := json.Marshal(b)
+	if err != nil {
+		return "", err
+	}
+	return string(data), nil
+}
+
+// sourceBoundFromJSON constructs SourceBound from its JSON represent.
+func sourceBoundFromJSON(s string) (b SourceBound, err error) {
+	err = json.Unmarshal([]byte(s), &b)
+	return
+}
+
+// PutSourceBound puts the bound relationship into etcd.
+// k/v: worker-name -> bound relationship.
+func PutSourceBound(cli *clientv3.Client, bound SourceBound) (int64, error) {
+	op, err := putSourceBoundOp(bound)
+	if err != nil {
+		return 0, err
+	}
+
+	ctx, cancel := context.WithTimeout(cli.Ctx(), etcdutil.DefaultRequestTimeout)
+	defer cancel()
+
+	resp, err := cli.Txn(ctx).Then(op).Commit()
+	if err != nil {
+		return 0, err
+	}
+	return resp.Header.Revision, nil
+}
+
+// GetSourceBound gets the source bound relationship for the specified DM-worker.
+func GetSourceBound(cli *clientv3.Client, worker string) (SourceBound, int64, error) {
+	ctx, cancel := context.WithTimeout(cli.Ctx(), etcdutil.DefaultRequestTimeout)
+	defer cancel()
+
+	bound := SourceBound{}
+	resp, err := cli.Get(ctx, common.UpstreamBoundWorkerKeyAdapter.Encode(worker))
+	if err != nil {
+		return bound, 0, err
+	}
+
+	if resp.Count == 0 {
+		return bound, 0, nil
+	} else if resp.Count > 1 {
+		// TODO(csuzhangxc): add terror.
+		// this should not happen.
+		return bound, 0, fmt.Errorf("too many bound relationship (%d) exist for the DM-worker %s", resp.Count, worker)
+	}
+
+	bound, err = sourceBoundFromJSON(string(resp.Kvs[0].Value))
+	if err != nil {
+		return bound, 0, err
+	}
+
+	return bound, resp.Header.Revision, nil
+}
+
+// WatchSourceBound watches PUT & DELETE operations for the bound relationship of the specified DM-worker.
+// For the DELETE operations, it returns an empty bound relationship.
+func WatchSourceBound(ctx context.Context, cli *clientv3.Client,
+	worker string, revision int64, outCh chan<- SourceBound) {
+	ch := cli.Watch(ctx, common.UpstreamBoundWorkerKeyAdapter.Encode(worker), clientv3.WithRev(revision))
+
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case resp := <-ch:
+			if resp.Canceled {
+				return
+			}
+
+			for _, ev := range resp.Events {
+				var (
+					bound SourceBound
+					err   error
+				)
+				switch ev.Type {
+				case mvccpb.PUT:
+					bound, err = sourceBoundFromJSON(string(ev.Kv.Value))
+					if err != nil {
+						// this should not happen.
+						log.L().Error("fail to construct source bound relationship", zap.ByteString("json", ev.Kv.Value))
+						continue
+					}
+				case mvccpb.DELETE:
+				default:
+					// this should not happen.
+					log.L().Error("unsupported etcd event type", zap.Reflect("kv", ev.Kv), zap.Reflect("type", ev.Type))
+					continue
+				}
+
+				select {
+				case outCh <- bound:
+				case <-ctx.Done():
+					return
+				}
+			}
+		}
+	}
+}
+
+// deleteSourceBoundOp returns a DELETE ectd operation for the bound relationship of the specified DM-worker.
+func deleteSourceBoundOp(worker string) clientv3.Op {
+	return clientv3.OpDelete(common.UpstreamBoundWorkerKeyAdapter.Encode(worker))
+}
+
+// putSourceBoundOp returns a PUT etcd operation for the bound relationship.
+// k/v: worker-name -> bound relationship.
+func putSourceBoundOp(bound SourceBound) (clientv3.Op, error) {
+	value, err := bound.toJSON()
+	if err != nil {
+		return clientv3.Op{}, err
+	}
+	key := common.UpstreamBoundWorkerKeyAdapter.Encode(bound.Worker)
+
+	return clientv3.OpPut(key, value), nil
+}

pkg/ha/bound_test.go

@@ -0,0 +1,95 @@
+// Copyright 2020 PingCAP, Inc.
+//
+// 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,
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package ha
+
+import (
+	"context"
+	"time"
+
+	. "github.com/pingcap/check"
+)
+
+func (t *testForEtcd) TestSourceBoundJSON(c *C) {
+	b1 := NewSourceBound("mysql-replica-1", "dm-worker-1")
+
+	j, err := b1.toJSON()
+	c.Assert(err, IsNil)
+	c.Assert(j, Equals, `{"source":"mysql-replica-1","worker":"dm-worker-1"}`)
+	c.Assert(j, Equals, b1.String())
+
+	b2, err := sourceBoundFromJSON(j)
+	c.Assert(err, IsNil)
+	c.Assert(b2, DeepEquals, b1)
+}
+
+func (t *testForEtcd) TestSourceBoundEtcd(c *C) {
+	defer clearTestInfoOperation(c)
+
+	var (
+		watchTimeout = 500 * time.Millisecond
+		worker       = "dm-worker-1"
+		emptyBound   = SourceBound{}
+		bound        = NewSourceBound("mysql-replica-1", worker)
+	)
+	c.Assert(bound.NotBound(), IsFalse)
+
+	// no bound exists.
+	bo1, rev1, err := GetSourceBound(etcdTestCli, worker)
+	c.Assert(err, IsNil)
+	c.Assert(rev1, Equals, int64(0))
+	c.Assert(bo1, DeepEquals, emptyBound)
+
+	// put a bound.
+	rev2, err := PutSourceBound(etcdTestCli, bound)
+	c.Assert(err, IsNil)
+	c.Assert(rev2, Greater, rev1)
+
+	// watch the PUT operation for the bound.
+	boundCh := make(chan SourceBound, 10)
+	ctx, cancel := context.WithTimeout(context.Background(), watchTimeout)
+	WatchSourceBound(ctx, etcdTestCli, worker, rev2, boundCh)
+	cancel()
+	close(boundCh)
+	c.Assert(len(boundCh), Equals, 1)
+	c.Assert(<-boundCh, DeepEquals, bound)
+
+	// get the bound back.
+	bo2, rev3, err := GetSourceBound(etcdTestCli, worker)
+	c.Assert(err, IsNil)
+	c.Assert(rev3, Equals, rev2)
+	c.Assert(bo2, DeepEquals, bound)
+
+	// delete the bound.
+	deleteOp := deleteSourceBoundOp(worker)
+	resp, err := etcdTestCli.Txn(context.Background()).Then(deleteOp).Commit()
+	c.Assert(err, IsNil)
+	rev4 := resp.Header.Revision
+	c.Assert(rev4, Greater, rev3)
+
+	// watch the DELETE operation for the bound.
+	boundCh = make(chan SourceBound, 10)
+	ctx, cancel = context.WithTimeout(context.Background(), watchTimeout)
+	WatchSourceBound(ctx, etcdTestCli, worker, rev4, boundCh)
+	cancel()
+	close(boundCh)
+	c.Assert(len(boundCh), Equals, 1)
+	bo3 := <-boundCh
+	c.Assert(bo3.NotBound(), IsTrue)
+
+	// get again, not exists now.
+	bo4, rev5, err := GetSourceBound(etcdTestCli, worker)
+	c.Assert(err, IsNil)
+	c.Assert(rev5, Equals, int64(0))
+	c.Assert(bo4, DeepEquals, emptyBound)
+}