package zookeeper import ( "context" "encoding/json" "errors" "fmt" "sync" "sync/atomic" "time" "github.com/bilibili/kratos/pkg/log" "github.com/bilibili/kratos/pkg/naming" "github.com/go-zookeeper/zk" ) type Config struct { // Endpoints is a list of URLs. Endpoints []string `json:"endpoints"` } var ( _once sync.Once _builder naming.Builder //ErrDuplication is a register duplication err ErrDuplication = errors.New("zookeeper: instance duplicate registration") ) // Builder return default zookeeper resolver builder. func Builder(c *Config) naming.Builder { _once.Do(func() { _builder, _ = New(c) }) return _builder } // Build register resolver into default zookeeper. func Build(c *Config, id string) naming.Resolver { return Builder(c).Build(id) } // ZookeeperBuilder is a zookeeper client Builder type ZookeeperBuilder struct { cli *zk.Conn connEvent <-chan zk.Event ctx context.Context cancelFunc context.CancelFunc mutex sync.RWMutex apps map[string]*appInfo registry map[string]struct{} } type appInfo struct { resolver map[*Resolve]struct{} ins atomic.Value zkb *ZookeeperBuilder once sync.Once } // Resolve zookeeper resolver. type Resolve struct { id string event chan struct{} zkb *ZookeeperBuilder } // New is new a zookeeper builder func New(c *Config) (zkb *ZookeeperBuilder, err error) { //example: endpointSli = []string{"192.168.1.78:2181", "192.168.1.79:2181", "192.168.1.80:2181"} if len(c.Endpoints) == 0 { errInfo := fmt.Sprintf("zookeeper New failed, endpoints is null") log.Error(errInfo) return nil, errors.New(errInfo) } zkConn, connEvent, err := zk.Connect(c.Endpoints, 5*time.Second) if err != nil { log.Error(fmt.Sprintf("zk Connect err:(%v)", err)) return } else { log.Info(fmt.Sprintf("zk Connect ok!")) } ctx, cancel := context.WithCancel(context.Background()) zkb = &ZookeeperBuilder{ cli: zkConn, connEvent: connEvent, ctx: ctx, cancelFunc: cancel, apps: map[string]*appInfo{}, registry: map[string]struct{}{}, } return } // Build zookeeper resovler builder. func (z *ZookeeperBuilder) Build(appid string) naming.Resolver { r := &Resolve{ id: appid, zkb: z, event: make(chan struct{}, 1), } z.mutex.Lock() app, ok := z.apps[appid] if !ok { app = &appInfo{ resolver: make(map[*Resolve]struct{}), zkb: z, } z.apps[appid] = app } app.resolver[r] = struct{}{} z.mutex.Unlock() if ok { select { case r.event <- struct{}{}: default: } } app.once.Do(func() { go app.watch(appid) log.Info("zookeeper: AddWatch(%s) already watch(%v)", appid, ok) }) return r } // Scheme return zookeeper's scheme func (z *ZookeeperBuilder) Scheme() string { return "zookeeper" } // Register is register instance func (z *ZookeeperBuilder) Register(ctx context.Context, ins *naming.Instance) (cancelFunc context.CancelFunc, err error) { z.mutex.Lock() if _, ok := z.registry[ins.AppID]; ok { err = ErrDuplication } else { z.registry[ins.AppID] = struct{}{} } z.mutex.Unlock() if err != nil { return } ctx, cancel := context.WithCancel(z.ctx) if err = z.register(ctx, ins); err != nil { z.mutex.Lock() delete(z.registry, ins.AppID) z.mutex.Unlock() cancel() return } ch := make(chan struct{}, 1) cancelFunc = context.CancelFunc(func() { cancel() <-ch }) go func() { for { select { case connEvent := <-z.connEvent: log.Warn("watch zkClient state, connEvent:(%v)", connEvent) if connEvent.State == zk.StateHasSession { log.Warn("watch zkClient state, state is StateHasSession...") err = z.register(ctx, ins) if err != nil { log.Warn(fmt.Sprintf("watch zkClient state, fail to register node error:(%v)", err)) continue } } case <-ctx.Done(): ch <- struct{}{} return } } }() return } func (z *ZookeeperBuilder) registerPerServer(name string) (err error) { var ( str string ) str, err = z.cli.Create(name, nil, 0, zk.WorldACL(zk.PermAll)) if err != nil { log.Warn(fmt.Sprintf("registerPerServer, fail to Create node:(%s). err:(%v)", name, err)) } else { log.Info(fmt.Sprintf("registerPerServer, succeed to Create node:(%s). retStr:(%s)", name, str)) } return } func (z *ZookeeperBuilder) registerEphServer(name, host string, ins *naming.Instance) (err error) { var ( str string ) val, _ := json.Marshal(ins) log.Info(fmt.Sprintf("registerEphServer, ins after json.Marshal:(%v)", string(val))) str, err = z.cli.Create(name+host, val, zk.FlagEphemeral, zk.WorldACL(zk.PermAll)) if err != nil { log.Warn(fmt.Sprintf("registerEphServer, fail to Create node:%s. err:(%v)", name+host, err)) } else { log.Info(fmt.Sprintf("registerEphServer, succeed to Create node:%s. retStr:(%s)", name+host, str)) } return } // register 注册zookeeper节点 func (z *ZookeeperBuilder) register(ctx context.Context, ins *naming.Instance) (err error) { log.Info("zookeeper register enter, instance Addrs:(%v)", ins.Addrs) prefix := z.keyPrefix(ins) err = z.registerPerServer(prefix) if err != nil { log.Warn(fmt.Sprintf("register, fail to registerPerServer node error:(%v)", err)) } for _, val := range ins.Addrs { err = z.registerEphServer(prefix, "/"+val, ins) if err != nil { log.Warn(fmt.Sprintf("registerServer, fail to RegisterEphServer node error:(%v)", err)) } else { log.Info(fmt.Sprintf("registerServer, succeed to RegistServer node.")) } } return nil } // unregister 删除zookeeper中节点信息 func (z *ZookeeperBuilder) unregister(ins *naming.Instance) (err error) { log.Info("zookeeper unregister enter, instance Addrs:(%v)", ins.Addrs) prefix := z.keyPrefix(ins) for _, val := range ins.Addrs { strNode := prefix + "/" + val exists, _, err := z.cli.Exists(strNode) if err != nil { log.Error("zk.Conn.Exists node:(%v), error:(%s)", strNode, err.Error()) return err } if exists { _, s, err := z.cli.Get(strNode) if err != nil { log.Error("zk.Conn.Get node:(%s), error:(%s)", strNode, err.Error()) return err } return z.cli.Delete(strNode, s.Version) } log.Info(fmt.Sprintf("unregister, client.Delete:(%v), appid:(%v), hostname:(%v) success", strNode, ins.AppID, ins.Hostname)) } return } func (z *ZookeeperBuilder) keyPrefix(ins *naming.Instance) string { return fmt.Sprintf("/%s", ins.AppID) } // Close stop all running process including zk fetch and register func (z *ZookeeperBuilder) Close() error { z.cancelFunc() return nil } func (a *appInfo) watch(appID string) { _ = a.fetchstore(appID) prefix := fmt.Sprintf("/%s", appID) go func() { for { log.Info(fmt.Sprintf("zk ChildrenW enter, prefix:(%v)", prefix)) snapshot, _, event, err := a.zkb.cli.ChildrenW(prefix) if err != nil { continue } log.Info(fmt.Sprintf("zk ChildrenW ok, snapshot:(%v)", snapshot)) for ev := range event { log.Info(fmt.Sprintf("zk ChildrenW ok, prefix:(%v), event Path:(%v), Type:(%v)", prefix, ev.Path, ev.Type)) if ev.Type == zk.EventNodeChildrenChanged { _ = a.fetchstore(appID) } } } }() } func (a *appInfo) fetchstore(appID string) (err error) { prefix := fmt.Sprintf("/%s", appID) strNode := "" childs, _, err := a.zkb.cli.Children(prefix) if err != nil { log.Error(fmt.Sprintf("fetchstore, fail to get Children of node:(%v), err:(%v)", prefix, err)) } else { log.Info(fmt.Sprintf("fetchstore, ok to get Children of node:(%v), childs:(%v)", prefix, childs)) } ins := &naming.InstancesInfo{ Instances: make(map[string][]*naming.Instance, 0), } //for range childs for _, child := range childs { strNode = prefix + "/" + child resp, _, err := a.zkb.cli.Get(strNode) if err != nil { log.Error("zookeeper: fetch client.Get(%s) error(%v)", strNode, err) return err } in := new(naming.Instance) err = json.Unmarshal(resp, in) if err != nil { return err } ins.Instances[in.Zone] = append(ins.Instances[in.Zone], in) } a.store(ins) return nil } func (a *appInfo) store(ins *naming.InstancesInfo) { a.ins.Store(ins) a.zkb.mutex.RLock() for rs := range a.resolver { select { case rs.event <- struct{}{}: default: } } a.zkb.mutex.RUnlock() } // Watch watch instance. func (r *Resolve) Watch() <-chan struct{} { return r.event } // Fetch fetch resolver instance. func (r *Resolve) Fetch(ctx context.Context) (ins *naming.InstancesInfo, ok bool) { r.zkb.mutex.RLock() app, ok := r.zkb.apps[r.id] r.zkb.mutex.RUnlock() if ok { ins, ok = app.ins.Load().(*naming.InstancesInfo) return } return } // Close close resolver. func (r *Resolve) Close() error { r.zkb.mutex.Lock() if app, ok := r.zkb.apps[r.id]; ok && len(app.resolver) != 0 { delete(app.resolver, r) } r.zkb.mutex.Unlock() return nil }