github.com/15mga/kiwi@v0.0.2-0.20240324021231-b95d5c3ac751/discovery/etcd/service.go (about)

     1  package etcd
     2  
     3  import (
     4  	"fmt"
     5  	"github.com/15mga/kiwi"
     6  	"github.com/15mga/kiwi/util"
     7  	"github.com/15mga/kiwi/util/etd"
     8  	"go.etcd.io/etcd/api/v3/mvccpb"
     9  	etcd "go.etcd.io/etcd/client/v3"
    10  	"strconv"
    11  	"strings"
    12  	"sync/atomic"
    13  )
    14  
    15  var (
    16  	RegTtl           int64 = 5
    17  	_RegRoot               = ""
    18  	_RegSvcIdxOffset       = 1
    19  	_SvcLeaseId      int64
    20  )
    21  
    22  func SetRegRoot(root string) {
    23  	_RegRoot = root
    24  	_RegSvcIdxOffset = strings.Count(RegSvcPrefix(), ".")
    25  }
    26  
    27  func RegSvcPrefix() string {
    28  	return _RegRoot + "node.info"
    29  }
    30  
    31  func RegLockPrefix() string {
    32  	return _RegRoot + "node.lock"
    33  }
    34  
    35  func RegisterService() {
    36  	//退出前注销
    37  	kiwi.BeforeExitFn("service unregister", unregisterSvc)
    38  	registerSvc()
    39  	go func() {
    40  		svcWatch := etd.Client().Watch(util.Ctx(), RegSvcPrefix(), etcd.WithPrefix())
    41  		for {
    42  			select {
    43  			case <-util.Ctx().Done():
    44  				return
    45  			case info := <-svcWatch:
    46  				rcvSvcEvent(info)
    47  			}
    48  		}
    49  	}()
    50  }
    51  
    52  func registerSvc() {
    53  	svcSlc := make([]*kiwi.NodeMeta, 0, 8)
    54  	err := etd.Lock(RegLockPrefix(), func() *util.Err {
    55  		nodeIdMap := make(map[int64]struct{}, 8)
    56  		_, e := etd.Get(RegSvcPrefix(), func(key string, bytes []byte) bool {
    57  			var si kiwi.NodeMeta
    58  			err := util.JsonUnmarshal(bytes, &si)
    59  			if err != nil {
    60  				return true
    61  			}
    62  			nodeIdMap[si.NodeId] = struct{}{}
    63  			svcSlc = append(svcSlc, &si)
    64  			return true
    65  		}, etcd.WithPrefix())
    66  		if e != nil {
    67  			return util.WrapErr(util.EcEtcdErr, e)
    68  		}
    69  		nodeId := int64(0)
    70  		info := kiwi.GetNodeMeta()
    71  		for ; nodeId < 1024; nodeId++ {
    72  			if _, ok := nodeIdMap[nodeId]; !ok {
    73  				info.SetSvcId(nodeId)
    74  				break
    75  			}
    76  		}
    77  		if nodeId == 0 {
    78  			return util.NewErr(util.EcServiceErr, util.M{
    79  				"error": "too much service node",
    80  			})
    81  		}
    82  
    83  		bytes, _ := util.JsonMarshal(info)
    84  		str := string(bytes)
    85  		key := getRegSvcKey(info.Svc, info.NodeId)
    86  		leaseId, err := etd.PutWithTtl(key, str, RegTtl)
    87  		if err != nil {
    88  			return err
    89  		}
    90  		atomic.StoreInt64(&_SvcLeaseId, leaseId)
    91  		kiwi.Info("register service success", util.M{
    92  			"node":         leaseId,
    93  			"lease id":     leaseId,
    94  			"service info": info,
    95  		})
    96  
    97  		return nil
    98  	})
    99  	if err != nil {
   100  		kiwi.Error(err)
   101  		return
   102  	}
   103  
   104  	for _, si := range svcSlc {
   105  		kiwi.Node().Connect(si.Ip, si.Port, si.Svc, si.SvcId, si.Ver, si.Data)
   106  	}
   107  }
   108  
   109  func unregisterSvc() {
   110  	id := atomic.SwapInt64(&_SvcLeaseId, 0)
   111  	if id == 0 {
   112  		return
   113  	}
   114  	kiwi.Info("unregister service", util.M{
   115  		"info": kiwi.GetNodeMeta(),
   116  	})
   117  	_ = etd.Revoke(id)
   118  }
   119  
   120  func rcvSvcEvent(res etcd.WatchResponse) {
   121  	for _, event := range res.Events {
   122  		switch event.Type {
   123  		case mvccpb.PUT:
   124  			var si kiwi.NodeMeta
   125  			err := util.JsonUnmarshal(event.Kv.Value, &si)
   126  			if err != nil {
   127  				kiwi.Error(err)
   128  				return
   129  			}
   130  			if si.SvcId == kiwi.GetNodeMeta().NodeId {
   131  				return
   132  			}
   133  			kiwi.Node().Connect(si.Ip, si.Port, si.Svc, si.SvcId, si.Ver, si.Data)
   134  		case mvccpb.DELETE:
   135  			key := string(event.Kv.Key)
   136  			svc, id, err := splitRegSvcKey(key)
   137  			if err != nil {
   138  				kiwi.Error(err)
   139  				return
   140  			}
   141  			if id == kiwi.GetNodeMeta().NodeId {
   142  				return
   143  			}
   144  			kiwi.Node().Disconnect(svc, id)
   145  		}
   146  	}
   147  }
   148  
   149  func getRegSvcKey(svc kiwi.TSvc, id int64) string {
   150  	return fmt.Sprintf("%s.%d.%d", RegSvcPrefix(), svc, id)
   151  }
   152  
   153  func splitRegSvcKey(key string) (svc kiwi.TSvc, id int64, err *util.Err) {
   154  	slc := strings.Split(key, ".")
   155  	l := len(slc)
   156  	if l != _RegSvcIdxOffset+3 {
   157  		err = util.NewErr(util.EcEtcdErr, util.M{
   158  			"key": key,
   159  		})
   160  		return
   161  	}
   162  	svci, e := strconv.Atoi(slc[_RegSvcIdxOffset+1])
   163  	if e != nil {
   164  		err = util.WrapErr(util.EcParseErr, e)
   165  		return
   166  	}
   167  	id, e = strconv.ParseInt(slc[_RegSvcIdxOffset+2], 10, 64)
   168  	if e != nil {
   169  		err = util.WrapErr(util.EcParseErr, e)
   170  		return
   171  	}
   172  	svc = kiwi.TSvc(svci)
   173  	return
   174  }