github.com/insionng/yougam@v0.0.0-20170714101924-2bc18d833463/libraries/ngaut/go-zookeeper/zk/conn.go (about) 1 package zk 2 3 /* 4 TODO: 5 * make sure a ping response comes back in a reasonable time 6 7 Possible watcher events: 8 * Event{Type: EventNotWatching, State: StateDisconnected, Path: path, Err: err} 9 */ 10 11 import ( 12 "crypto/rand" 13 "encoding/binary" 14 "errors" 15 "fmt" 16 "io" 17 "log" 18 "net" 19 "strconv" 20 "strings" 21 "sync" 22 "sync/atomic" 23 "time" 24 ) 25 26 var ErrNoServer = errors.New("zk: could not connect to a server") 27 28 const ( 29 bufferSize = 10 * 1024 * 1024 30 eventChanSize = 6 31 sendChanSize = 16 32 protectedPrefix = "_c_" 33 ) 34 35 type watchType int 36 37 const ( 38 watchTypeData = iota 39 watchTypeExist = iota 40 watchTypeChild = iota 41 ) 42 43 type watchPathType struct { 44 path string 45 wType watchType 46 } 47 48 type Dialer func(network, address string, timeout time.Duration) (net.Conn, error) 49 50 type Conn struct { 51 lastZxid int64 52 sessionID int64 53 state State // must be 32-bit aligned 54 xid int32 55 timeout int32 // session timeout in seconds 56 passwd []byte 57 58 dialer Dialer 59 servers []string 60 serverIndex int 61 conn net.Conn 62 eventChan chan Event 63 shouldQuit chan bool 64 pingInterval time.Duration 65 recvTimeout time.Duration 66 connectTimeout time.Duration 67 68 sendChan chan *request 69 requests map[int32]*request // Xid -> pending request 70 requestsLock sync.Mutex 71 watchers map[watchPathType][]chan Event 72 watchersLock sync.Mutex 73 74 // Debug (used by unit tests) 75 reconnectDelay time.Duration 76 } 77 78 type request struct { 79 xid int32 80 opcode int32 81 pkt interface{} 82 recvStruct interface{} 83 recvChan chan response 84 85 // Because sending and receiving happen in separate go routines, there's 86 // a possible race condition when creating watches from outside the read 87 // loop. We must ensure that a watcher gets added to the list synchronously 88 // with the response from the server on any request that creates a watch. 89 // In order to not hard code the watch logic for each opcode in the recv 90 // loop the caller can use recvFunc to insert some synchronously code 91 // after a response. 92 recvFunc func(*request, *responseHeader, error) 93 } 94 95 type response struct { 96 zxid int64 97 err error 98 } 99 100 type Event struct { 101 Type EventType 102 State State 103 Path string // For non-session events, the path of the watched node. 104 Err error 105 } 106 107 func Connect(servers []string, recvTimeout time.Duration) (*Conn, <-chan Event, error) { 108 return ConnectWithDialer(servers, recvTimeout, nil) 109 } 110 111 func ConnectWithDialer(servers []string, recvTimeout time.Duration, dialer Dialer) (*Conn, <-chan Event, error) { 112 // Randomize the order of the servers to avoid creating hotspots 113 stringShuffle(servers) 114 115 for i, addr := range servers { 116 if !strings.Contains(addr, ":") { 117 servers[i] = addr + ":" + strconv.Itoa(DefaultPort) 118 } 119 } 120 ec := make(chan Event, eventChanSize) 121 if dialer == nil { 122 dialer = net.DialTimeout 123 } 124 conn := Conn{ 125 dialer: dialer, 126 servers: servers, 127 serverIndex: 0, 128 conn: nil, 129 state: StateDisconnected, 130 eventChan: ec, 131 shouldQuit: make(chan bool), 132 recvTimeout: recvTimeout, 133 pingInterval: time.Duration((int64(recvTimeout) / 2)), 134 connectTimeout: 1 * time.Second, 135 sendChan: make(chan *request, sendChanSize), 136 requests: make(map[int32]*request), 137 watchers: make(map[watchPathType][]chan Event), 138 passwd: emptyPassword, 139 timeout: 30000, 140 141 // Debug 142 reconnectDelay: time.Second, 143 } 144 go func() { 145 conn.loop() 146 conn.flushRequests(ErrClosing) 147 conn.invalidateWatches(ErrClosing) 148 close(conn.eventChan) 149 }() 150 return &conn, ec, nil 151 } 152 153 func (c *Conn) Close() { 154 close(c.shouldQuit) 155 156 select { 157 case <-c.queueRequest(opClose, &closeRequest{}, &closeResponse{}, nil): 158 case <-time.After(time.Second): 159 } 160 } 161 162 func (c *Conn) State() State { 163 return State(atomic.LoadInt32((*int32)(&c.state))) 164 } 165 166 func (c *Conn) setState(state State) { 167 atomic.StoreInt32((*int32)(&c.state), int32(state)) 168 select { 169 case c.eventChan <- Event{Type: EventSession, State: state}: 170 default: 171 // panic("zk: event channel full - it must be monitored and never allowed to be full") 172 } 173 } 174 175 func (c *Conn) connect() { 176 c.serverIndex = (c.serverIndex + 1) % len(c.servers) 177 startIndex := c.serverIndex 178 c.setState(StateConnecting) 179 for { 180 zkConn, err := c.dialer("tcp", c.servers[c.serverIndex], c.connectTimeout) 181 if err == nil { 182 c.conn = zkConn 183 c.setState(StateConnected) 184 return 185 } 186 187 log.Printf("Failed to connect to %s: %+v", c.servers[c.serverIndex], err) 188 189 c.serverIndex = (c.serverIndex + 1) % len(c.servers) 190 if c.serverIndex == startIndex { 191 c.flushUnsentRequests(ErrNoServer) 192 time.Sleep(time.Second) 193 } 194 } 195 } 196 197 func (c *Conn) loop() { 198 for { 199 c.connect() 200 err := c.authenticate() 201 switch { 202 case err == ErrSessionExpired: 203 c.invalidateWatches(err) 204 case err != nil && c.conn != nil: 205 c.conn.Close() 206 case err == nil: 207 closeChan := make(chan bool) // channel to tell send loop stop 208 var wg sync.WaitGroup 209 210 wg.Add(1) 211 go func() { 212 c.sendLoop(c.conn, closeChan) 213 c.conn.Close() // causes recv loop to EOF/exit 214 wg.Done() 215 }() 216 217 wg.Add(1) 218 go func() { 219 err = c.recvLoop(c.conn) 220 if err == nil { 221 panic("zk: recvLoop should never return nil error") 222 } 223 close(closeChan) // tell send loop to exit 224 wg.Done() 225 }() 226 227 wg.Wait() 228 } 229 230 c.setState(StateDisconnected) 231 232 // Yeesh 233 if err != io.EOF && err != ErrSessionExpired && !strings.Contains(err.Error(), "use of closed network connection") { 234 log.Println(err) 235 } 236 237 select { 238 case <-c.shouldQuit: 239 c.flushRequests(ErrClosing) 240 return 241 default: 242 } 243 244 if err != ErrSessionExpired { 245 err = ErrConnectionClosed 246 } 247 c.flushRequests(err) 248 249 if c.reconnectDelay > 0 { 250 select { 251 case <-c.shouldQuit: 252 return 253 case <-time.After(c.reconnectDelay): 254 } 255 } 256 } 257 } 258 259 func (c *Conn) flushUnsentRequests(err error) { 260 for { 261 select { 262 default: 263 return 264 case req := <-c.sendChan: 265 req.recvChan <- response{-1, err} 266 } 267 } 268 } 269 270 // Send error to all pending requests and clear request map 271 func (c *Conn) flushRequests(err error) { 272 c.requestsLock.Lock() 273 for _, req := range c.requests { 274 req.recvChan <- response{-1, err} 275 } 276 c.requests = make(map[int32]*request) 277 c.requestsLock.Unlock() 278 } 279 280 // Send error to all watchers and clear watchers map 281 func (c *Conn) invalidateWatches(err error) { 282 c.watchersLock.Lock() 283 defer c.watchersLock.Unlock() 284 285 if len(c.watchers) >= 0 { 286 for pathType, watchers := range c.watchers { 287 ev := Event{Type: EventNotWatching, State: StateDisconnected, Path: pathType.path, Err: err} 288 for _, ch := range watchers { 289 ch <- ev 290 close(ch) 291 } 292 } 293 c.watchers = make(map[watchPathType][]chan Event) 294 } 295 } 296 297 func (c *Conn) sendSetWatches() { 298 c.watchersLock.Lock() 299 defer c.watchersLock.Unlock() 300 301 if len(c.watchers) == 0 { 302 return 303 } 304 305 req := &setWatchesRequest{ 306 RelativeZxid: c.lastZxid, 307 DataWatches: make([]string, 0), 308 ExistWatches: make([]string, 0), 309 ChildWatches: make([]string, 0), 310 } 311 n := 0 312 for pathType, watchers := range c.watchers { 313 if len(watchers) == 0 { 314 continue 315 } 316 switch pathType.wType { 317 case watchTypeData: 318 req.DataWatches = append(req.DataWatches, pathType.path) 319 case watchTypeExist: 320 req.ExistWatches = append(req.ExistWatches, pathType.path) 321 case watchTypeChild: 322 req.ChildWatches = append(req.ChildWatches, pathType.path) 323 } 324 n++ 325 } 326 if n == 0 { 327 return 328 } 329 330 go func() { 331 res := &setWatchesResponse{} 332 _, err := c.request(opSetWatches, req, res, nil) 333 if err != nil { 334 log.Printf("Failed to set previous watches: %s", err.Error()) 335 } 336 }() 337 } 338 339 func (c *Conn) authenticate() error { 340 buf := make([]byte, 256) 341 342 // connect request 343 344 n, err := encodePacket(buf[4:], &connectRequest{ 345 ProtocolVersion: protocolVersion, 346 LastZxidSeen: c.lastZxid, 347 TimeOut: c.timeout, 348 SessionID: c.sessionID, 349 Passwd: c.passwd, 350 }) 351 if err != nil { 352 return err 353 } 354 355 binary.BigEndian.PutUint32(buf[:4], uint32(n)) 356 357 _, err = c.conn.Write(buf[:n+4]) 358 if err != nil { 359 return err 360 } 361 362 c.sendSetWatches() 363 364 // connect response 365 366 // package length 367 _, err = io.ReadFull(c.conn, buf[:4]) 368 if err != nil { 369 return err 370 } 371 372 blen := int(binary.BigEndian.Uint32(buf[:4])) 373 if cap(buf) < blen { 374 buf = make([]byte, blen) 375 } 376 377 _, err = io.ReadFull(c.conn, buf[:blen]) 378 if err != nil { 379 return err 380 } 381 382 r := connectResponse{} 383 _, err = decodePacket(buf[:blen], &r) 384 if err != nil { 385 return err 386 } 387 if r.SessionID == 0 { 388 c.sessionID = 0 389 c.passwd = emptyPassword 390 c.lastZxid = 0 391 c.setState(StateExpired) 392 return ErrSessionExpired 393 } 394 395 if c.sessionID != r.SessionID { 396 atomic.StoreInt32(&c.xid, 0) 397 } 398 c.timeout = r.TimeOut 399 c.sessionID = r.SessionID 400 c.passwd = r.Passwd 401 c.setState(StateHasSession) 402 403 return nil 404 } 405 406 func (c *Conn) sendLoop(conn net.Conn, closeChan <-chan bool) error { 407 pingTicker := time.NewTicker(c.pingInterval) 408 defer pingTicker.Stop() 409 410 buf := make([]byte, bufferSize) 411 for { 412 select { 413 case req := <-c.sendChan: 414 header := &requestHeader{req.xid, req.opcode} 415 n, err := encodePacket(buf[4:], header) 416 if err != nil { 417 req.recvChan <- response{-1, err} 418 continue 419 } 420 421 n2, err := encodePacket(buf[4+n:], req.pkt) 422 if err != nil { 423 req.recvChan <- response{-1, err} 424 continue 425 } 426 427 n += n2 428 429 binary.BigEndian.PutUint32(buf[:4], uint32(n)) 430 431 c.requestsLock.Lock() 432 select { 433 case <-closeChan: 434 req.recvChan <- response{-1, ErrConnectionClosed} 435 c.requestsLock.Unlock() 436 return ErrConnectionClosed 437 default: 438 } 439 c.requests[req.xid] = req 440 c.requestsLock.Unlock() 441 442 conn.SetWriteDeadline(time.Now().Add(c.recvTimeout)) 443 _, err = conn.Write(buf[:n+4]) 444 conn.SetWriteDeadline(time.Time{}) 445 if err != nil { 446 req.recvChan <- response{-1, err} 447 conn.Close() 448 return err 449 } 450 case <-pingTicker.C: 451 n, err := encodePacket(buf[4:], &requestHeader{Xid: -2, Opcode: opPing}) 452 if err != nil { 453 panic("zk: opPing should never fail to serialize") 454 } 455 456 binary.BigEndian.PutUint32(buf[:4], uint32(n)) 457 458 conn.SetWriteDeadline(time.Now().Add(c.recvTimeout)) 459 _, err = conn.Write(buf[:n+4]) 460 conn.SetWriteDeadline(time.Time{}) 461 if err != nil { 462 conn.Close() 463 return err 464 } 465 case <-closeChan: 466 return nil 467 } 468 } 469 } 470 471 func (c *Conn) recvLoop(conn net.Conn) error { 472 buf := make([]byte, bufferSize) 473 for { 474 // package length 475 conn.SetReadDeadline(time.Now().Add(c.recvTimeout)) 476 _, err := io.ReadFull(conn, buf[:4]) 477 if err != nil { 478 return err 479 } 480 481 blen := int(binary.BigEndian.Uint32(buf[:4])) 482 if cap(buf) < blen { 483 buf = make([]byte, blen) 484 } 485 486 _, err = io.ReadFull(conn, buf[:blen]) 487 conn.SetReadDeadline(time.Time{}) 488 if err != nil { 489 return err 490 } 491 492 res := responseHeader{} 493 _, err = decodePacket(buf[:16], &res) 494 if err != nil { 495 return err 496 } 497 498 if res.Xid == -1 { 499 res := &watcherEvent{} 500 _, err := decodePacket(buf[16:16+blen], res) 501 if err != nil { 502 return err 503 } 504 ev := Event{ 505 Type: res.Type, 506 State: res.State, 507 Path: res.Path, 508 Err: nil, 509 } 510 select { 511 case c.eventChan <- ev: 512 default: 513 } 514 wTypes := make([]watchType, 0, 2) 515 switch res.Type { 516 case EventNodeCreated: 517 wTypes = append(wTypes, watchTypeExist) 518 case EventNodeDeleted, EventNodeDataChanged: 519 wTypes = append(wTypes, watchTypeExist, watchTypeData, watchTypeChild) 520 case EventNodeChildrenChanged: 521 wTypes = append(wTypes, watchTypeChild) 522 } 523 c.watchersLock.Lock() 524 for _, t := range wTypes { 525 wpt := watchPathType{res.Path, t} 526 if watchers := c.watchers[wpt]; watchers != nil && len(watchers) > 0 { 527 for _, ch := range watchers { 528 ch <- ev 529 close(ch) 530 } 531 delete(c.watchers, wpt) 532 } 533 } 534 c.watchersLock.Unlock() 535 } else if res.Xid == -2 { 536 // Ping response. Ignore. 537 } else if res.Xid < 0 { 538 log.Printf("Xid < 0 (%d) but not ping or watcher event", res.Xid) 539 } else { 540 if res.Zxid > 0 { 541 c.lastZxid = res.Zxid 542 } 543 544 c.requestsLock.Lock() 545 req, ok := c.requests[res.Xid] 546 if ok { 547 delete(c.requests, res.Xid) 548 } 549 c.requestsLock.Unlock() 550 551 if !ok { 552 log.Printf("Response for unknown request with xid %d", res.Xid) 553 } else { 554 if res.Err != 0 { 555 err = res.Err.toError() 556 } else { 557 _, err = decodePacket(buf[16:16+blen], req.recvStruct) 558 } 559 if req.recvFunc != nil { 560 req.recvFunc(req, &res, err) 561 } 562 req.recvChan <- response{res.Zxid, err} 563 if req.opcode == opClose { 564 return io.EOF 565 } 566 } 567 } 568 } 569 } 570 571 func (c *Conn) nextXid() int32 { 572 return atomic.AddInt32(&c.xid, 1) 573 } 574 575 func (c *Conn) addWatcher(path string, watchType watchType) <-chan Event { 576 c.watchersLock.Lock() 577 defer c.watchersLock.Unlock() 578 579 ch := make(chan Event, 1) 580 wpt := watchPathType{path, watchType} 581 c.watchers[wpt] = append(c.watchers[wpt], ch) 582 return ch 583 } 584 585 func (c *Conn) queueRequest(opcode int32, req interface{}, res interface{}, recvFunc func(*request, *responseHeader, error)) <-chan response { 586 rq := &request{ 587 xid: c.nextXid(), 588 opcode: opcode, 589 pkt: req, 590 recvStruct: res, 591 recvChan: make(chan response, 1), 592 recvFunc: recvFunc, 593 } 594 c.sendChan <- rq 595 return rq.recvChan 596 } 597 598 func (c *Conn) request(opcode int32, req interface{}, res interface{}, recvFunc func(*request, *responseHeader, error)) (int64, error) { 599 r := <-c.queueRequest(opcode, req, res, recvFunc) 600 return r.zxid, r.err 601 } 602 603 func (c *Conn) AddAuth(scheme string, auth []byte) error { 604 _, err := c.request(opSetAuth, &setAuthRequest{Type: 0, Scheme: scheme, Auth: auth}, &setAuthResponse{}, nil) 605 return err 606 } 607 608 func (c *Conn) Children(path string) ([]string, Stat, error) { 609 res := &getChildren2Response{} 610 _, err := c.request(opGetChildren2, &getChildren2Request{Path: path, Watch: false}, res, nil) 611 return res.Children, &res.Stat, err 612 } 613 614 func (c *Conn) ChildrenW(path string) ([]string, Stat, <-chan Event, error) { 615 var ech <-chan Event 616 res := &getChildren2Response{} 617 _, err := c.request(opGetChildren2, &getChildren2Request{Path: path, Watch: true}, res, func(req *request, res *responseHeader, err error) { 618 if err == nil { 619 ech = c.addWatcher(path, watchTypeChild) 620 } 621 }) 622 if err != nil { 623 return nil, nil, nil, err 624 } 625 return res.Children, &res.Stat, ech, err 626 } 627 628 func (c *Conn) Get(path string) ([]byte, Stat, error) { 629 res := &getDataResponse{} 630 _, err := c.request(opGetData, &getDataRequest{Path: path, Watch: false}, res, nil) 631 return res.Data, &res.Stat, err 632 } 633 634 // GetW returns the contents of a znode and sets a watch 635 func (c *Conn) GetW(path string) ([]byte, Stat, <-chan Event, error) { 636 var ech <-chan Event 637 res := &getDataResponse{} 638 _, err := c.request(opGetData, &getDataRequest{Path: path, Watch: true}, res, func(req *request, res *responseHeader, err error) { 639 if err == nil { 640 ech = c.addWatcher(path, watchTypeData) 641 } 642 }) 643 if err != nil { 644 return nil, nil, nil, err 645 } 646 return res.Data, &res.Stat, ech, err 647 } 648 649 func (c *Conn) Set(path string, data []byte, version int32) (Stat, error) { 650 res := &setDataResponse{} 651 _, err := c.request(opSetData, &SetDataRequest{path, data, version}, res, nil) 652 return &res.Stat, err 653 } 654 655 func (c *Conn) Create(path string, data []byte, flags int32, acl []ACL) (string, error) { 656 res := &createResponse{} 657 _, err := c.request(opCreate, &CreateRequest{path, data, acl, flags}, res, nil) 658 return res.Path, err 659 } 660 661 // CreateProtectedEphemeralSequential fixes a race condition if the server crashes 662 // after it creates the node. On reconnect the session may still be valid so the 663 // ephemeral node still exists. Therefore, on reconnect we need to check if a node 664 // with a GUID generated on create exists. 665 func (c *Conn) CreateProtectedEphemeralSequential(path string, data []byte, acl []ACL) (string, error) { 666 var guid [16]byte 667 _, err := io.ReadFull(rand.Reader, guid[:16]) 668 if err != nil { 669 return "", err 670 } 671 guidStr := fmt.Sprintf("%x", guid) 672 673 parts := strings.Split(path, "/") 674 parts[len(parts)-1] = fmt.Sprintf("%s%s-%s", protectedPrefix, guidStr, parts[len(parts)-1]) 675 rootPath := strings.Join(parts[:len(parts)-1], "/") 676 protectedPath := strings.Join(parts, "/") 677 678 var newPath string 679 for i := 0; i < 3; i++ { 680 newPath, err = c.Create(protectedPath, data, FlagEphemeral|FlagSequence, acl) 681 switch err { 682 case ErrSessionExpired: 683 // No need to search for the node since it can't exist. Just try again. 684 case ErrConnectionClosed: 685 children, _, err := c.Children(rootPath) 686 if err != nil { 687 return "", err 688 } 689 for _, p := range children { 690 parts := strings.Split(p, "/") 691 if pth := parts[len(parts)-1]; strings.HasPrefix(pth, protectedPrefix) { 692 if g := pth[len(protectedPrefix) : len(protectedPrefix)+32]; g == guidStr { 693 return rootPath + "/" + p, nil 694 } 695 } 696 } 697 case nil: 698 return newPath, nil 699 default: 700 return "", err 701 } 702 } 703 return "", err 704 } 705 706 func (c *Conn) Delete(path string, version int32) error { 707 _, err := c.request(opDelete, &DeleteRequest{path, version}, &deleteResponse{}, nil) 708 return err 709 } 710 711 func (c *Conn) Exists(path string) (bool, Stat, error) { 712 res := &existsResponse{} 713 _, err := c.request(opExists, &existsRequest{Path: path, Watch: false}, res, nil) 714 exists := true 715 if err == ErrNoNode { 716 exists = false 717 err = nil 718 } 719 return exists, &res.Stat, err 720 } 721 722 func (c *Conn) ExistsW(path string) (bool, Stat, <-chan Event, error) { 723 var ech <-chan Event 724 res := &existsResponse{} 725 _, err := c.request(opExists, &existsRequest{Path: path, Watch: true}, res, func(req *request, res *responseHeader, err error) { 726 if err == nil { 727 ech = c.addWatcher(path, watchTypeData) 728 } else if err == ErrNoNode { 729 ech = c.addWatcher(path, watchTypeExist) 730 } 731 }) 732 exists := true 733 if err == ErrNoNode { 734 exists = false 735 err = nil 736 } 737 if err != nil { 738 return false, nil, nil, err 739 } 740 return exists, &res.Stat, ech, err 741 } 742 743 func (c *Conn) GetACL(path string) ([]ACL, Stat, error) { 744 res := &getAclResponse{} 745 _, err := c.request(opGetAcl, &getAclRequest{Path: path}, res, nil) 746 return res.Acl, &res.Stat, err 747 } 748 749 func (c *Conn) SetACL(path string, acl []ACL, version int32) (Stat, error) { 750 res := &setAclResponse{} 751 _, err := c.request(opSetAcl, &setAclRequest{Path: path, Acl: acl, Version: version}, res, nil) 752 return &res.Stat, err 753 } 754 755 func (c *Conn) Sync(path string) (string, error) { 756 res := &syncResponse{} 757 _, err := c.request(opSync, &syncRequest{Path: path}, res, nil) 758 return res.Path, err 759 } 760 761 type MultiOps struct { 762 Create []CreateRequest 763 Delete []DeleteRequest 764 SetData []SetDataRequest 765 Check []CheckVersionRequest 766 } 767 768 func (c *Conn) Multi(ops MultiOps) error { 769 req := &multiRequest{ 770 Ops: make([]multiRequestOp, 0, len(ops.Create)+len(ops.Delete)+len(ops.SetData)+len(ops.Check)), 771 DoneHeader: multiHeader{Type: -1, Done: true, Err: -1}, 772 } 773 for _, r := range ops.Create { 774 req.Ops = append(req.Ops, multiRequestOp{multiHeader{opCreate, false, -1}, r}) 775 } 776 for _, r := range ops.SetData { 777 req.Ops = append(req.Ops, multiRequestOp{multiHeader{opSetData, false, -1}, r}) 778 } 779 for _, r := range ops.Delete { 780 req.Ops = append(req.Ops, multiRequestOp{multiHeader{opDelete, false, -1}, r}) 781 } 782 for _, r := range ops.Check { 783 req.Ops = append(req.Ops, multiRequestOp{multiHeader{opCheck, false, -1}, r}) 784 } 785 res := &multiResponse{} 786 _, err := c.request(opMulti, req, res, nil) 787 return err 788 }