github.com/go-kratos/kratos@v1.0.1/pkg/naming/zookeeper/zookeeper.go (about) 1 package zookeeper 2 3 import ( 4 "context" 5 "encoding/json" 6 "errors" 7 "fmt" 8 "net/url" 9 "path" 10 "strings" 11 "sync" 12 "sync/atomic" 13 "time" 14 15 "github.com/go-zookeeper/zk" 16 17 "github.com/go-kratos/kratos/pkg/log" 18 "github.com/go-kratos/kratos/pkg/naming" 19 xtime "github.com/go-kratos/kratos/pkg/time" 20 ) 21 22 // Config is zookeeper config. 23 type Config struct { 24 Root string `json:"root"` 25 Endpoints []string `json:"endpoints"` 26 Timeout xtime.Duration `json:"timeout"` 27 } 28 29 var ( 30 _once sync.Once 31 _builder naming.Builder 32 33 // ErrDuplication is a register duplication err 34 ErrDuplication = errors.New("zookeeper: instance duplicate registration") 35 ) 36 37 // Builder return default zookeeper resolver builder. 38 func Builder(c *Config) naming.Builder { 39 _once.Do(func() { 40 _builder, _ = New(c) 41 }) 42 return _builder 43 } 44 45 // Build register resolver into default zookeeper. 46 func Build(c *Config, id string) naming.Resolver { 47 return Builder(c).Build(id) 48 } 49 50 type appInfo struct { 51 resolver map[*Resolve]struct{} 52 ins atomic.Value 53 zkb *Zookeeper 54 once sync.Once 55 } 56 57 // Resolve zookeeper resolver. 58 type Resolve struct { 59 id string 60 event chan struct{} 61 zkb *Zookeeper 62 } 63 64 // Zookeeper is a zookeeper client Builder. 65 // path: /{root}/{appid}/{ip} -> json(instance) 66 type Zookeeper struct { 67 c *Config 68 cli *zk.Conn 69 connEvent <-chan zk.Event 70 ctx context.Context 71 cancelFunc context.CancelFunc 72 73 mutex sync.RWMutex 74 apps map[string]*appInfo 75 registry map[string]struct{} 76 } 77 78 // New is new a zookeeper builder. 79 func New(c *Config) (zkb *Zookeeper, err error) { 80 if c.Timeout == 0 { 81 c.Timeout = xtime.Duration(time.Second) 82 } 83 if len(c.Endpoints) == 0 { 84 errInfo := "zookeeper New failed, endpoints is null" 85 log.Error(errInfo) 86 return nil, errors.New(errInfo) 87 } 88 89 zkConn, connEvent, err := zk.Connect(c.Endpoints, time.Duration(c.Timeout)) 90 if err != nil { 91 log.Error(fmt.Sprintf("zk Connect err:(%v)", err)) 92 return 93 } 94 log.Info("zk Connect ok!") 95 96 ctx, cancel := context.WithCancel(context.Background()) 97 zkb = &Zookeeper{ 98 c: c, 99 cli: zkConn, 100 connEvent: connEvent, 101 ctx: ctx, 102 cancelFunc: cancel, 103 apps: map[string]*appInfo{}, 104 registry: map[string]struct{}{}, 105 } 106 return 107 } 108 109 // Build zookeeper resovler builder. 110 func (z *Zookeeper) Build(appid string, options ...naming.BuildOpt) naming.Resolver { 111 r := &Resolve{ 112 id: appid, 113 zkb: z, 114 event: make(chan struct{}, 1), 115 } 116 z.mutex.Lock() 117 app, ok := z.apps[appid] 118 if !ok { 119 app = &appInfo{ 120 resolver: make(map[*Resolve]struct{}), 121 zkb: z, 122 } 123 z.apps[appid] = app 124 } 125 app.resolver[r] = struct{}{} 126 z.mutex.Unlock() 127 if ok { 128 select { 129 case r.event <- struct{}{}: 130 default: 131 } 132 } 133 app.once.Do(func() { 134 go app.watch(appid) 135 }) 136 return r 137 } 138 139 // Scheme return zookeeper's scheme. 140 func (z *Zookeeper) Scheme() string { 141 return "zookeeper" 142 } 143 144 // Register is register instance. 145 func (z *Zookeeper) Register(ctx context.Context, ins *naming.Instance) (cancelFunc context.CancelFunc, err error) { 146 z.mutex.Lock() 147 if _, ok := z.registry[ins.AppID]; ok { 148 err = ErrDuplication 149 } else { 150 z.registry[ins.AppID] = struct{}{} 151 } 152 z.mutex.Unlock() 153 if err != nil { 154 return 155 } 156 ctx, cancel := context.WithCancel(z.ctx) 157 if err = z.register(ctx, ins); err != nil { 158 z.mutex.Lock() 159 delete(z.registry, ins.AppID) 160 z.mutex.Unlock() 161 cancel() 162 return 163 } 164 ch := make(chan struct{}, 1) 165 cancelFunc = context.CancelFunc(func() { 166 cancel() 167 <-ch 168 }) 169 go func() { 170 for { 171 select { 172 case connEvent := <-z.connEvent: 173 log.Info("watch zkClient state, connEvent:(%+v)", connEvent) 174 if connEvent.State == zk.StateHasSession { 175 if err = z.register(ctx, ins); err != nil { 176 log.Warn(fmt.Sprintf("watch zkClient state, fail to register node error:(%v)", err)) 177 continue 178 } 179 } 180 case <-ctx.Done(): 181 ch <- struct{}{} 182 return 183 } 184 } 185 }() 186 return 187 } 188 189 func (z *Zookeeper) createPath(paths string) error { 190 var ( 191 lastPath = "/" 192 seps = strings.Split(paths, "/") 193 ) 194 for _, part := range seps { 195 if part == "" { 196 continue 197 } 198 lastPath = path.Join(lastPath, part) 199 ok, _, err := z.cli.Exists(lastPath) 200 if err != nil { 201 return err 202 } 203 if ok { 204 continue 205 } 206 ret, err := z.cli.Create(lastPath, nil, 0, zk.WorldACL(zk.PermAll)) 207 if err != nil { 208 log.Warn(fmt.Sprintf("createPath, fail to Create node:(%s). error:(%v)", paths, err)) 209 } else { 210 log.Info(fmt.Sprintf("createPath, succeed to Create node:(%s). retStr:(%s)", paths, ret)) 211 } 212 } 213 return nil 214 } 215 216 func (z *Zookeeper) registerPeerServer(nodePath string, ins *naming.Instance) (err error) { 217 var ( 218 str string 219 ) 220 val, err := json.Marshal(ins) 221 if err != nil { 222 return 223 } 224 log.Info(fmt.Sprintf("registerPeerServer, ins after json.Marshal:(%v)", string(val))) 225 ok, _, err := z.cli.Exists(nodePath) 226 if err != nil { 227 return err 228 } 229 if ok { 230 return nil 231 } 232 str, err = z.cli.Create(nodePath, val, zk.FlagEphemeral, zk.WorldACL(zk.PermAll)) 233 if err != nil { 234 log.Warn(fmt.Sprintf("registerPeerServer, fail to Create node:%s. error:(%v)", nodePath, err)) 235 } else { 236 log.Info(fmt.Sprintf("registerPeerServer, succeed to Create node:%s. retStr:(%s)", nodePath, str)) 237 } 238 return 239 } 240 241 // register is register instance to zookeeper. 242 func (z *Zookeeper) register(ctx context.Context, ins *naming.Instance) (err error) { 243 log.Info("zookeeper register enter, instance Addrs:(%v)", ins.Addrs) 244 245 prefix := z.keyPrefix(ins.AppID) 246 if err = z.createPath(prefix); err != nil { 247 log.Warn(fmt.Sprintf("register, fail to createPath node error:(%v)", err)) 248 } 249 for _, addr := range ins.Addrs { 250 u, err := url.Parse(addr) 251 if err != nil { 252 continue 253 } 254 // grpc://127.0.0.1:8000 to 127.0.0.1 255 nodePath := prefix + "/" + strings.SplitN(u.Host, ":", 2)[0] 256 if err = z.registerPeerServer(nodePath, ins); err != nil { 257 log.Warn(fmt.Sprintf("registerServer, fail to RegisterPeerServer node:%s error:(%v)", addr, err)) 258 } else { 259 log.Info("registerServer, succeed to RegistServer node.") 260 } 261 } 262 return nil 263 } 264 265 func (z *Zookeeper) unregister(ins *naming.Instance) (err error) { 266 log.Info("zookeeper unregister enter, instance Addrs:(%v)", ins.Addrs) 267 prefix := z.keyPrefix(ins.AppID) 268 for _, addr := range ins.Addrs { 269 u, err := url.Parse(addr) 270 if err != nil { 271 continue 272 } 273 // grpc://127.0.0.1:8000 to 127.0.0.1 274 nodePath := prefix + "/" + strings.SplitN(u.Host, ":", 2)[0] 275 exists, _, err := z.cli.Exists(nodePath) 276 if err != nil { 277 log.Error("zk.Conn.Exists node:(%v), error:(%v)", nodePath, err) 278 continue 279 } 280 if exists { 281 _, s, err := z.cli.Get(nodePath) 282 if err != nil { 283 log.Error("zk.Conn.Get node:(%s), error:(%v)", nodePath, err) 284 continue 285 } 286 if err = z.cli.Delete(nodePath, s.Version); err != nil { 287 log.Error("zk.Conn.Delete node:(%s), error:(%v)", nodePath, err) 288 continue 289 } 290 } 291 292 log.Info(fmt.Sprintf("unregister, client.Delete:(%v), appid:(%v), hostname:(%v) success", nodePath, ins.AppID, ins.Hostname)) 293 } 294 return 295 } 296 297 func (z *Zookeeper) keyPrefix(appID string) string { 298 return path.Join(z.c.Root, appID) 299 } 300 301 // Close stop all running process including zk fetch and register. 302 func (z *Zookeeper) Close() error { 303 z.cancelFunc() 304 return nil 305 } 306 307 func (a *appInfo) watch(appID string) { 308 _ = a.fetchstore(appID) 309 go func() { 310 prefix := a.zkb.keyPrefix(appID) 311 for { 312 log.Info(fmt.Sprintf("zk ChildrenW enter, prefix:(%v)", prefix)) 313 snapshot, _, event, err := a.zkb.cli.ChildrenW(prefix) 314 if err != nil { 315 log.Error("zk ChildrenW fail to watch:%s error:(%v)", prefix, err) 316 time.Sleep(time.Second) 317 _ = a.fetchstore(appID) 318 continue 319 } 320 log.Info(fmt.Sprintf("zk ChildrenW ok, prefix:%s snapshot:(%v)", prefix, snapshot)) 321 for ev := range event { 322 log.Info(fmt.Sprintf("zk ChildrenW ok, prefix:(%v), event Path:(%v), Type:(%v)", prefix, ev.Path, ev.Type)) 323 if ev.Type == zk.EventNodeChildrenChanged { 324 _ = a.fetchstore(appID) 325 } 326 } 327 } 328 }() 329 } 330 331 func (a *appInfo) fetchstore(appID string) (err error) { 332 prefix := a.zkb.keyPrefix(appID) 333 childs, _, err := a.zkb.cli.Children(prefix) 334 if err != nil { 335 log.Error(fmt.Sprintf("fetchstore, fail to get Children of node:(%v), error:(%v)", prefix, err)) 336 return 337 } 338 log.Info(fmt.Sprintf("fetchstore, ok to get Children of node:(%v), childs:(%v)", prefix, childs)) 339 ins := &naming.InstancesInfo{ 340 Instances: make(map[string][]*naming.Instance), 341 } 342 for _, child := range childs { 343 nodePath := prefix + "/" + child 344 resp, _, err := a.zkb.cli.Get(nodePath) 345 if err != nil { 346 log.Error("zookeeper: fetch client.Get(%s) error:(%v)", nodePath, err) 347 return err 348 } 349 in := new(naming.Instance) 350 if err = json.Unmarshal(resp, in); err != nil { 351 return err 352 } 353 ins.Instances[in.Zone] = append(ins.Instances[in.Zone], in) 354 } 355 a.store(ins) 356 return nil 357 } 358 359 func (a *appInfo) store(ins *naming.InstancesInfo) { 360 a.ins.Store(ins) 361 a.zkb.mutex.RLock() 362 for rs := range a.resolver { 363 select { 364 case rs.event <- struct{}{}: 365 default: 366 } 367 } 368 a.zkb.mutex.RUnlock() 369 } 370 371 // Watch watch instance. 372 func (r *Resolve) Watch() <-chan struct{} { 373 return r.event 374 } 375 376 // Fetch fetch resolver instance. 377 func (r *Resolve) Fetch(ctx context.Context) (ins *naming.InstancesInfo, ok bool) { 378 r.zkb.mutex.RLock() 379 app, ok := r.zkb.apps[r.id] 380 r.zkb.mutex.RUnlock() 381 if ok { 382 ins, ok = app.ins.Load().(*naming.InstancesInfo) 383 return 384 } 385 return 386 } 387 388 // Close close resolver. 389 func (r *Resolve) Close() error { 390 r.zkb.mutex.Lock() 391 if app, ok := r.zkb.apps[r.id]; ok && len(app.resolver) != 0 { 392 delete(app.resolver, r) 393 } 394 r.zkb.mutex.Unlock() 395 return nil 396 }