// hanleTCPConn handle a long live tcp connection. func handleTCPConn(conn net.Conn, rc chan *bufio.Reader) { addr := conn.RemoteAddr().String() log.Debug("<%s> handleTcpConn routine start", addr) rd := newBufioReader(rc, conn) if args, err := parseCmd(rd); err == nil { // return buffer bufio.Reader putBufioReader(rc, rd) switch args[0] { case "sub": SubscribeTCPHandle(conn, args[1:]) break default: conn.Write(ParamReply) log.Warn("<%s> unknown cmd \"%s\"", addr, args[0]) break } } else { // return buffer bufio.Reader putBufioReader(rc, rd) log.Error("<%s> parseCmd() error(%v)", addr, err) } // close the connection if err := conn.Close(); err != nil { log.Error("<%s> conn.Close() error(%v)", addr, err) } log.Debug("<%s> handleTcpConn routine stop", addr) return }
// HandleWrite start a goroutine get msg from chan, then send to the conn. func (c *Connection) HandleWrite(key string) { go func() { var ( n int err error ) log.Debug("user_key: \"%s\" HandleWrite goroutine start", key) for { msg, ok := <-c.Buf if !ok { log.Debug("user_key: \"%s\" HandleWrite goroutine stop", key) return } if c.Proto == WebsocketProto { // raw n, err = c.Conn.Write(msg) } else if c.Proto == TCPProto { // redis protocol msg = []byte(fmt.Sprintf("$%d\r\n%s\r\n", len(msg), string(msg))) n, err = c.Conn.Write(msg) } else { log.Error("unknown connection protocol: %d", c.Proto) panic(ErrConnProto) } // update stat if err != nil { log.Error("user_key: \"%s\" conn.Write() error(%v)", key, err) MsgStat.IncrFailed(1) } else { log.Debug("user_key: \"%s\" write \r\n========%s(%d)========", key, string(msg), n) MsgStat.IncrSucceed(1) } } }() }
// Migrate migrate portion of connections which don't belong to this comet. func (l *ChannelList) Migrate(nw map[string]int) (err error) { migrate := false // check new/update node for k, v := range nw { weight, ok := nodeWeightMap[k] // not found or weight change if !ok || weight != v { migrate = true break } } // check del node if !migrate { for k, _ := range nodeWeightMap { // node deleted if _, ok := nw[k]; !ok { migrate = true break } } } if !migrate { return } // init ketama ring := ketama.NewRing(ketama.Base) for node, weight := range nw { ring.AddNode(node, weight) } ring.Bake() // atomic update nodeWeightMap = nw CometRing = ring // get all the channel lock channels := []Channel{} for i, c := range l.Channels { c.Lock() for k, v := range c.Data { hn := ring.Hash(k) if hn != Conf.ZookeeperCometNode { channels = append(channels, v) delete(c.Data, k) log.Debug("migrate delete channel key \"%s\"", k) } } c.Unlock() log.Debug("migrate channel bucket:%d finished", i) } // close all the migrate channels log.Info("close all the migrate channels") for _, channel := range channels { if err := channel.Close(); err != nil { log.Error("channel.Close() error(%v)", err) continue } } log.Info("close all the migrate channels finished") return }
func tcpListen(bind string) { addr, err := net.ResolveTCPAddr("tcp", bind) if err != nil { log.Error("net.ResolveTCPAddr(\"tcp\"), %s) error(%v)", bind, err) panic(err) } l, err := net.ListenTCP("tcp", addr) if err != nil { log.Error("net.ListenTCP(\"tcp4\", \"%s\") error(%v)", bind, err) panic(err) } // free the listener resource defer func() { log.Info("tcp addr: \"%s\" close", bind) if err := l.Close(); err != nil { log.Error("listener.Close() error(%v)", err) } }() // init reader buffer instance rb := newtcpBufCache() for { log.Debug("start accept") conn, err := l.AcceptTCP() if err != nil { log.Error("listener.AcceptTCP() error(%v)", err) continue } if err = conn.SetKeepAlive(Conf.TCPKeepalive); err != nil { log.Error("conn.SetKeepAlive() error(%v)", err) conn.Close() continue } if err = conn.SetReadBuffer(Conf.RcvbufSize); err != nil { log.Error("conn.SetReadBuffer(%d) error(%v)", Conf.RcvbufSize, err) conn.Close() continue } if err = conn.SetWriteBuffer(Conf.SndbufSize); err != nil { log.Error("conn.SetWriteBuffer(%d) error(%v)", Conf.SndbufSize, err) conn.Close() continue } // first packet must sent by client in specified seconds if err = conn.SetReadDeadline(time.Now().Add(fitstPacketTimedoutSec)); err != nil { log.Error("conn.SetReadDeadLine() error(%v)", err) conn.Close() continue } rc := rb.Get() // one connection one routine go handleTCPConn(conn, rc) log.Debug("accept finished") } }
// validate check the key is belong to this comet. func (l *ChannelList) validate(key string) error { if len(nodeWeightMap) == 0 { log.Debug("no node found") return ErrChannelKey } node := CometRing.Hash(key) log.Debug("match node:%s hash node:%s", Conf.ZookeeperCometNode, node) if Conf.ZookeeperCometNode != node { log.Warn("user_key:\"%s\" node:%s not match this node:%s", key, node, Conf.ZookeeperCometNode) return ErrChannelKey } return nil }
// GetPrivate rpc interface get user private message. func (r *MessageRPC) GetPrivate(m *myrpc.MessageGetPrivateArgs, rw *myrpc.MessageGetResp) error { log.Debug("messageRPC.GetPrivate key:\"%s\" mid:\"%d\"", m.Key, m.MsgId) if m == nil || m.Key == "" || m.MsgId < 0 { return myrpc.ErrParam } msgs, err := UseStorage.GetPrivate(m.Key, m.MsgId) if err != nil { log.Error("UseStorage.GetPrivate(\"%s\", %d) error(%v)", m.Key, m.MsgId, err) return err } rw.Msgs = msgs log.Debug("UserStorage.GetPrivate(\"%s\", %d) ok", m.Key, m.MsgId) return nil }
// RegisterTmp create a ephemeral node, and watch it, if node droped then send a SIGQUIT to self. func RegisterTemp(conn *zk.Conn, fpath string, data []byte) error { tpath, err := conn.Create(path.Join(fpath)+"/", data, zk.FlagEphemeral|zk.FlagSequence, zk.WorldACL(zk.PermAll)) if err != nil { log.Error("conn.Create(\"%s\", \"%s\", zk.FlagEphemeral|zk.FlagSequence) error(%v)", fpath, string(data), err) return err } log.Debug("create a zookeeper node:%s", tpath) // watch self go func() { for { log.Info("zk path: \"%s\" set a watch", tpath) exist, _, watch, err := conn.ExistsW(tpath) if err != nil { log.Error("zk.ExistsW(\"%s\") error(%v)", tpath, err) log.Warn("zk path: \"%s\" set watch failed, kill itself", tpath) killSelf() return } if !exist { log.Warn("zk path: \"%s\" not exist, kill itself", tpath) killSelf() return } event := <-watch log.Info("zk path: \"%s\" receive a event %v", tpath, event) } }() return nil }
// bucket return a channelBucket use murmurhash3. func (l *ChannelList) Bucket(key string) *ChannelBucket { h := hash.NewMurmur3C() h.Write([]byte(key)) idx := uint(h.Sum32()) & uint(Conf.ChannelBucket-1) log.Debug("user_key:\"%s\" hit channel bucket index:%d", key, idx) return l.Channels[idx] }
// newTCPBufCache return a new tcpBuf cache. func newtcpBufCache() *tcpBufCache { inst := make([]chan *bufio.Reader, 0, Conf.BufioInstance) log.Debug("create %d read buffer instance", Conf.BufioInstance) for i := 0; i < Conf.BufioInstance; i++ { inst = append(inst, make(chan *bufio.Reader, Conf.BufioNum)) } return &tcpBufCache{instance: inst, round: 0} }
// getConn get the connection of matching with key using ketama hashing. func (s *RedisStorage) getConn(key string) redis.Conn { if len(s.pool) == 0 { return nil } node := s.ring.Hash(key) log.Debug("user_key: \"%s\" hit redis node: \"%s\"", key, node) return s.getConnByNode(node) }
// GetComet get the node infomation under the node. func GetComet(key string) *CometNodeInfo { if cometRing == nil || len(cometNodeInfoMap) == 0 { return nil } node := cometRing.Hash(key) log.Debug("cometHash hits \"%s\"", node) return cometNodeInfoMap[node] }
// SavePrivate rpc interface save user private message. func (r *MessageRPC) SavePrivate(m *myrpc.MessageSavePrivateArgs, ret *int) error { if m == nil || m.Msg == nil || m.MsgId < 0 { return myrpc.ErrParam } if err := UseStorage.SavePrivate(m.Key, m.Msg, m.MsgId, m.Expire); err != nil { log.Error("UseStorage.SavePrivate(\"%s\", \"%s\", %d, %d) error(%v)", m.Key, string(m.Msg), m.MsgId, m.Expire, err) return err } log.Debug("UseStorage.SavePrivate(\"%s\", \"%s\", %d, %d) ok", m.Key, string(m.Msg), m.MsgId, m.Expire) return nil }
// notify every Comet node to migrate func notifyMigrate(conn *zk.Conn, migrateLockPath, znode, key string, update bool, nodeWeightMap map[string]int) (err error) { // try lock if _, err = conn.Create(migrateLockPath, []byte("1"), zk.FlagEphemeral, zk.WorldACL(zk.PermAll)); err != nil { log.Error("conn.Create(\"/gopush-migrate-lock\", \"1\", zk.FlagEphemeral) error(%v)", err) return } // call comet migrate rpc wg := &sync.WaitGroup{} wg.Add(len(cometNodeInfoMap)) for node, nodeInfo := range cometNodeInfoMap { go func(n string, info *CometNodeInfo) { if info.Rpc == nil { log.Error("notify migrate failed, no rpc found, node:%s", n) wg.Done() return } r := info.Rpc.Get() if r == nil { log.Error("notify migrate failed, no rpc found, node:%s", n) wg.Done() return } reply := 0 args := &CometMigrateArgs{Nodes: nodeWeightMap} if err = r.Call(CometServiceMigrate, args, &reply); err != nil { log.Error("rpc.Call(\"%s\") error(%v)", CometServiceMigrate, err) wg.Done() return } log.Debug("notify node:%s migrate succeed", n) wg.Done() }(node, nodeInfo) } wg.Wait() // update znode info if update { var data []byte data, err = json.Marshal(cometNodeInfoMap[key]) if err != nil { log.Error("json.Marshal() node:%s error(%v)", key, err) return } if _, err = conn.Set(znode, data, -1); err != nil { log.Error("conn.Set(\"%s\",\"%s\",\"-1\") error(%v)", znode, string(data), err) return } } // release lock if err = conn.Delete(migrateLockPath, -1); err != nil { log.Error("conn.Delete(\"%s\") error(%v)", migrateLockPath, err) } return }
// DelPrivate rpc interface delete user private message. func (r *MessageRPC) DelPrivate(key string, ret *int) error { if key == "" { return myrpc.ErrParam } if err := UseStorage.DelPrivate(key); err != nil { log.Error("UserStorage.DelPrivate(\"%s\") error(%v)", key, err) return err } log.Debug("UserStorage.DelPrivate(\"%s\") ok", key) return nil }
// watchMessageRoot watch the message root path. func watchMessageRoot(conn *zk.Conn, fpath string, ch chan *MessageNodeEvent) error { for { nodes, watch, err := myzk.GetNodesW(conn, fpath) if err == myzk.ErrNodeNotExist { log.Warn("zk don't have node \"%s\", retry in %d second", fpath, waitNodeDelay) time.Sleep(waitNodeDelaySecond) continue } else if err == myzk.ErrNoChild { log.Warn("zk don't have any children in \"%s\", retry in %d second", fpath, waitNodeDelay) // all child died, kick all the nodes for _, client := range MessageRPC.Clients { log.Debug("node: \"%s\" send del node event", client.Addr) ch <- &MessageNodeEvent{Event: eventNodeDel, Key: &WeightRpc{Addr: client.Addr, Weight: client.Weight}} } time.Sleep(waitNodeDelaySecond) continue } else if err != nil { log.Error("getNodes error(%v), retry in %d second", err, waitNodeDelay) time.Sleep(waitNodeDelaySecond) continue } nodesMap := map[string]bool{} // handle new add nodes for _, node := range nodes { data, _, err := conn.Get(path.Join(fpath, node)) if err != nil { log.Error("zk.Get(\"%s\") error(%v)", path.Join(fpath, node), err) continue } // parse message node info nodeInfo := &MessageNodeInfo{} if err := json.Unmarshal(data, nodeInfo); err != nil { log.Error("json.Unmarshal(\"%s\", nodeInfo) error(%v)", string(data), err) continue } for _, addr := range nodeInfo.Rpc { // if not exists in old map then trigger a add event if _, ok := MessageRPC.Clients[addr]; !ok { ch <- &MessageNodeEvent{Event: eventNodeAdd, Key: &WeightRpc{Addr: addr, Weight: nodeInfo.Weight}} } nodesMap[addr] = true } } // handle delete nodes for _, client := range MessageRPC.Clients { if _, ok := nodesMap[client.Addr]; !ok { ch <- &MessageNodeEvent{Event: eventNodeDel, Key: client} } } // blocking wait node changed event := <-watch log.Info("zk path: \"%s\" receive a event %v", fpath, event) } }
// SavePrivates rpc interface save user private messages. func (r *MessageRPC) SavePrivates(m *myrpc.MessageSavePrivatesArgs, rw *myrpc.MessageSavePrivatesResp) error { if m == nil || m.Msg == nil || m.MsgId < 0 { return myrpc.ErrParam } fkeys, err := UseStorage.SavePrivates(m.Keys, m.Msg, m.MsgId, m.Expire) if err != nil { log.Error("UseStorage.SavePrivates(\"%v\", \"%s\", %d, %d) error(%v)", m.Keys, string(m.Msg), m.MsgId, m.Expire, err) } rw.FKeys = fkeys log.Debug("UseStorage.SavePrivates(\"%v\", \"%s\", %d, %d) ok fkeys len(%d)", m.Keys, string(m.Msg), m.MsgId, m.Expire, len(fkeys)) return nil }
// registerCometNode get infomation of comet node func registerCometNode(conn *zk.Conn, node, fpath string, retry, ping time.Duration, startPing bool) (info *CometNodeInfo, err error) { // get current node info from zookeeper fpath = path.Join(fpath, node) data, _, err := conn.Get(fpath) if err != nil { log.Error("zk.Get(\"%s\") error(%v)", fpath, err) return } info = &CometNodeInfo{} if err = json.Unmarshal(data, info); err != nil { log.Error("json.Unmarshal(\"%s\", nodeData) error(%v)", string(data), err) return } if len(info.RpcAddr) == 0 { log.Error("zk nodes: \"%s\" don't have rpc addr", fpath) err = ErrCometRPC return } // get old node info for finding the old rpc connection oldInfo := cometNodeInfoMap[node] // init comet rpc clients := make(map[string]*WeightRpc, len(info.RpcAddr)) for _, addr := range info.RpcAddr { var ( r *rpc.Client ) if oldInfo != nil && oldInfo.Rpc != nil { if wr, ok := oldInfo.Rpc.Clients[addr]; ok && wr.Client != nil { // reuse the rpc connection must let old client = nil, avoid reclose rpc. oldInfo.Rpc.Clients[addr].Client = nil r = wr.Client } } if r == nil { if r, err = rpc.Dial("tcp", addr); err != nil { log.Error("rpc.Dial(\"%s\") error(%v)", addr, err) return } log.Debug("node:%s addr:%s rpc reconnect", node, addr) } clients[addr] = &WeightRpc{Weight: 1, Addr: addr, Client: r} } // comet rpc use rand load balance lb, err := NewRandLB(clients, cometService, retry, ping, startPing) if err != nil { log.Error("NewRandLR() error(%v)", err) return } info.Rpc = lb log.Info("zk path: \"%s\" register nodes: \"%s\"", fpath, node) return }
// NewChannelList create a new channel bucket set. func NewChannelList() *ChannelList { l := &ChannelList{Channels: []*ChannelBucket{}} // split hashmap to many bucket log.Debug("create %d ChannelBucket", Conf.ChannelBucket) for i := 0; i < Conf.ChannelBucket; i++ { c := &ChannelBucket{ Data: map[string]Channel{}, mutex: &sync.Mutex{}, } l.Channels = append(l.Channels, c) } return l }
// retPWrite marshal the result and write to client(post). func retPWrite(w http.ResponseWriter, r *http.Request, res map[string]interface{}, body *string, start time.Time) { data, err := json.Marshal(res) if err != nil { log.Error("json.Marshal(\"%v\") error(%v)", res, err) return } dataStr := string(data) if n, err := w.Write([]byte(dataStr)); err != nil { log.Error("w.Write(\"%s\") error(%v)", dataStr, err) } else { log.Debug("w.Write(\"%s\") write %d bytes", dataStr, n) } log.Info("req: \"%s\", post: \"%s\", res:\"%s\", ip:\"%s\", time:\"%fs\"", r.URL.String(), *body, dataStr, r.RemoteAddr, time.Now().Sub(start).Seconds()) }
// Connect connect to zookeeper, and start a goroutine log the event. func Connect(addr []string, timeout time.Duration) (*zk.Conn, error) { conn, session, err := zk.Connect(addr, timeout) if err != nil { log.Error("zk.Connect(\"%v\", %d) error(%v)", addr, timeout, err) return nil, err } go func() { for { event := <-session log.Debug("zookeeper get a event: %s", event.State.String()) } }() return conn, nil }
// Create create zookeeper path, if path exists ignore error func Create(conn *zk.Conn, fpath string) error { // create zk root path tpath := "" for _, str := range strings.Split(fpath, "/")[1:] { tpath = path.Join(tpath, "/", str) log.Debug("create zookeeper path: \"%s\"", tpath) _, err := conn.Create(tpath, []byte(""), 0, zk.WorldACL(zk.PermAll)) if err != nil { if err == zk.ErrNodeExists { log.Warn("zk.create(\"%s\") exists", tpath) } else { log.Error("zk.create(\"%s\") error(%v)", tpath, err) return err } } } return nil }
// handleNodeEvent add and remove MessageRPC.Clients, copy the src map to a new map then replace the variable. func handleMessageNodeEvent(conn *zk.Conn, retry, ping time.Duration, ch chan *MessageNodeEvent) { for { ev := <-ch // copy map from src tmpMessageRPCMap := make(map[string]*WeightRpc, len(MessageRPC.Clients)) for k, v := range MessageRPC.Clients { tmpMessageRPCMap[k] = &WeightRpc{Client: v.Client, Addr: v.Addr, Weight: v.Weight} // reuse rpc connection v.Client = nil } // handle event if ev.Event == eventNodeAdd { log.Info("add message rpc node: \"%s\"", ev.Key.Addr) rpcTmp, err := rpc.Dial("tcp", ev.Key.Addr) if err != nil { log.Error("rpc.Dial(\"tcp\", \"%s\") error(%v)", ev.Key, err) log.Warn("discard message rpc node: \"%s\", connect failed", ev.Key) continue } ev.Key.Client = rpcTmp tmpMessageRPCMap[ev.Key.Addr] = ev.Key } else if ev.Event == eventNodeDel { log.Info("del message rpc node: \"%s\"", ev.Key.Addr) delete(tmpMessageRPCMap, ev.Key.Addr) } else { log.Error("unknown node event: %d", ev.Event) panic("unknown node event") } tmpMessageRPC, err := NewRandLB(tmpMessageRPCMap, MessageService, retry, ping, true) if err != nil { log.Error("NewRandLR() error(%v)", err) panic(err) } oldMessageRPC := MessageRPC // atomic update MessageRPC = tmpMessageRPC // release resource oldMessageRPC.Destroy() log.Debug("MessageRPC.Client length: %d", len(MessageRPC.Clients)) } }
// retWrite marshal the result and write to client(get). func retWrite(w http.ResponseWriter, r *http.Request, res map[string]interface{}, callback string, start time.Time) { data, err := json.Marshal(res) if err != nil { log.Error("json.Marshal(\"%v\") error(%v)", res, err) return } dataStr := "" if callback == "" { // Normal json dataStr = string(data) } else { // Jsonp dataStr = fmt.Sprintf("%s(%s)", callback, string(data)) } if n, err := w.Write([]byte(dataStr)); err != nil { log.Error("w.Write(\"%s\") error(%v)", dataStr, err) } else { log.Debug("w.Write(\"%s\") write %d bytes", dataStr, n) } log.Info("req: \"%s\", res:\"%s\", ip:\"%s\", time:\"%fs\"", r.URL.String(), dataStr, r.RemoteAddr, time.Now().Sub(start).Seconds()) }
// Ping check health. func (c *CometRPC) Ping(args int, ret *int) error { log.Debug("ping ok") return nil }
// PushPrivates expored a method for publishing a user multiple private message for the channel. // because of it`s going asynchronously in this method, so it won`t return an error to caller. func (c *CometRPC) PushPrivates(args *myrpc.CometPushPrivatesArgs, rw *myrpc.CometPushPrivatesResp) error { if args == nil { return myrpc.ErrParam } bucketMap := make(map[*ChannelBucket]*batchChannel, Conf.ChannelBucket) for _, key := range args.Keys { // get channel ch, bp, err := UserChannel.New(key) if err != nil { log.Error("UserChannel.New(\"%s\") error(%v)", key, err) // log failed keys. rw.FKeys = append(rw.FKeys, key) continue } if bucket, ok := bucketMap[bp]; !ok { bucketMap[bp] = &batchChannel{ Keys: []string{key}, Chs: map[string]Channel{key: ch}, } } else { // ignore duplicate key if _, ok := bucket.Chs[key]; !ok { bucket.Chs[key] = ch bucket.Keys = append(bucket.Keys, key) } } } // every bucket start a goroutine, return till all bucket gorouint finish wg := &sync.WaitGroup{} wg.Add(len(bucketMap)) // stored every gorouint failed keys fKeysList := make([][]string, len(bucketMap)) ti := 0 for tb, tm := range bucketMap { go func(b *ChannelBucket, m *batchChannel, i int) { defer wg.Done() c := myrpc.MessageRPC.Get() if c == nil { // static slice is thread-safe // log all keys fKeysList[i] = m.Keys log.Debug("fkeys len:%d", len(m.Keys)) return } b.Lock() defer b.Unlock() timeId := id.Get() msg := &myrpc.Message{Msg: args.Msg, MsgId: timeId} // private message need persistence // if message expired no need persistence, only send online message // rewrite message id resp := &myrpc.MessageSavePrivatesResp{} if args.Expire > 0 { args := &myrpc.MessageSavePrivatesArgs{Keys: m.Keys, Msg: args.Msg, MsgId: timeId, Expire: args.Expire} if err := c.Call(myrpc.MessageServiceSavePrivates, args, resp); err != nil { log.Error("%s(\"%v\", \"%v\", &ret) error(%v)", myrpc.MessageServiceSavePrivates, m.Keys, args, err) // static slice is thread-safe fKeysList[i] = m.Keys log.Debug("fkeys len:%d", len(m.Keys)) return } fKeysList[i] = resp.FKeys log.Debug("fkeys len:%d", len(resp.FKeys)) } // delete the failed keys for _, fk := range resp.FKeys { delete(m.Chs, fk) } // get all channels from batchChannel chs. for key, ch := range m.Chs { if err := ch.WriteMsg(key, msg); err != nil { // ignore online push error, cause offline msg succeed log.Error("ch.WriteMsg(\"%s\", \"%s\") error(%v)", key, string(msg.Msg), err) continue } } }(tb, tm, ti) ti++ } wg.Wait() // merge all failed keys for _, k := range fKeysList { rw.FKeys = append(rw.FKeys, k...) } return nil }
// PushMultiPrivate handle for push multiple private messages. // Because of it`s going asynchronously in this method, so it won`t return a InternalErr to caller. func PushMultiPrivate(w http.ResponseWriter, r *http.Request) { if r.Method != "POST" { http.Error(w, "Method Not Allowed", 405) return } body := "" res := map[string]interface{}{"ret": OK} defer retPWrite(w, r, res, &body, time.Now()) // post param bodyBytes, err := ioutil.ReadAll(r.Body) if err != nil { res["ret"] = InternalErr log.Error("ioutil.ReadAll() failed (%v)", err) return } body = string(bodyBytes) msgBytes, keys, ret := parseMultiPrivate(bodyBytes) if ret != OK { res["ret"] = ret return } rm := json.RawMessage(msgBytes) msg, err := rm.MarshalJSON() if err != nil { res["ret"] = ParamErr log.Error("json.RawMessage(\"%s\").MarshalJSON() error(%v)", string(msg), err) return } // url param params := r.URL.Query() expire, err := strconv.ParseUint(params.Get("expire"), 10, 32) if err != nil { res["ret"] = ParamErr log.Error("strconv.ParseUint(\"%s\", 10, 32) error(%v)", params.Get("expire"), err) return } // match nodes nodes := map[*myrpc.CometNodeInfo]*[]string{} for i := 0; i < len(keys); i++ { node := myrpc.GetComet(keys[i]) if node == nil || node.Rpc == nil { res["ret"] = NotFoundServer return } keysTmp, ok := nodes[node] if ok { *keysTmp = append(*keysTmp, keys[i]) } else { nodes[node] = &([]string{keys[i]}) } } var fKeys []string //push to every node for cometInfo, ks := range nodes { client := cometInfo.Rpc.Get() if client == nil { log.Error("cannot get comet rpc client") fKeys = append(fKeys, *ks...) continue } args := &myrpc.CometPushPrivatesArgs{Msg: json.RawMessage(msg), Expire: uint(expire), Keys: *ks} resp := myrpc.CometPushPrivatesResp{} if err := client.Call(myrpc.CometServicePushPrivates, args, &resp); err != nil { log.Error("client.Call(\"%s\", \"%v\", &ret) error(%v)", myrpc.CometServicePushPrivates, args.Keys, err) fKeys = append(fKeys, *ks...) continue } log.Debug("fkeys len(%d) addr:%v", len(resp.FKeys), cometInfo.RpcAddr) fKeys = append(fKeys, resp.FKeys...) } res["ret"] = OK if len(fKeys) != 0 { res["data"] = map[string]interface{}{"fk": fKeys} } return }
// handleCometNodeEvent add and remove CometNodeInfo, copy the src map to a new map then replace the variable. func handleCometNodeEvent(conn *zk.Conn, migrateLockPath, fpath string, retry, ping time.Duration, ch chan *CometNodeEvent) { for { ev := <-ch var ( update = false znode = path.Join(fpath, ev.Key) ) // copy map from src tmpMap := make(map[string]*CometNodeInfo, len(cometNodeInfoMap)) for k, v := range cometNodeInfoMap { tmpMap[k] = v } // handle event if ev.Event == eventNodeAdd { log.Info("add node: \"%s\"", ev.Key) tmpMap[ev.Key] = &CometNodeInfo{Weight: 1} go watchCometNode(conn, ev.Key, fpath, retry, ping, ch) } else if ev.Event == eventNodeDel { log.Info("del node: \"%s\"", ev.Key) delete(tmpMap, ev.Key) } else if ev.Event == eventNodeUpdate { log.Info("update node: \"%s\"", ev.Key) // when new node added to watchCometNode then trigger node update tmpMap[ev.Key] = ev.Value update = true } else { log.Error("unknown node event: %d", ev.Event) panic("unknown node event") } // if exist old node info, destroy // if node add this may not happan // if node del this will clean the resource // if node update, after reuse rpc connection, this will clean the resource if info, ok := cometNodeInfoMap[ev.Key]; ok { if info != nil && info.Rpc != nil { info.Rpc.Destroy() } } // update comet hash, cause node has changed tempRing := ketama.NewRing(ketama.Base) nodeWeightMap := map[string]int{} for k, v := range tmpMap { log.Debug("AddNode node:%s weight:%d", k, v.Weight) tempRing.AddNode(k, v.Weight) nodeWeightMap[k] = v.Weight } tempRing.Bake() // use the tmpMap atomic replace the global cometNodeInfoMap cometNodeInfoMap = tmpMap cometRing = tempRing // migrate if ev.Event != eventNodeAdd { if err := notifyMigrate(conn, migrateLockPath, znode, ev.Key, update, nodeWeightMap); err != nil { // if err == zk.ErrNodeExists meaning anyone is going through. // we hopefully that only one web node notify comet migrate. // also it was judged in Comet whether it needs migrate or not. if err == zk.ErrNodeExists { log.Info("ignore notify migrate") continue } else { log.Error("notifyMigrate(conn, \"%v\") error(%v)", nodeWeightMap, err) continue } } } log.Debug("cometNodeInfoMap len: %d", len(cometNodeInfoMap)) } }
// Subscriber Handle is the websocket handle for sub request. func SubscribeHandle(ws *websocket.Conn) { addr := ws.Request().RemoteAddr params := ws.Request().URL.Query() // get subscriber key key := params.Get("key") if key == "" { ws.Write(ParamReply) log.Warn("<%s> key param error", addr) return } // get heartbeat second heartbeatStr := params.Get("heartbeat") i, err := strconv.Atoi(heartbeatStr) if err != nil { ws.Write(ParamReply) log.Error("<%s> user_key:\"%s\" heartbeat argument error(%v)", addr, key, err) return } if i < minHearbeatSec { ws.Write(ParamReply) log.Warn("<%s> user_key:\"%s\" heartbeat argument error, less than %d", addr, key, minHearbeatSec) return } heartbeat := i + delayHeartbeatSec token := params.Get("token") version := params.Get("ver") log.Info("<%s> subscribe to key = %s, heartbeat = %d, token = %s, version = %s", addr, key, heartbeat, token, version) // fetch subscriber from the channel c, err := UserChannel.Get(key, true) if err != nil { log.Warn("<%s> user_key:\"%s\" can't get a channel (%s)", addr, key, err) if err == ErrChannelKey { ws.Write(NodeReply) } else { ws.Write(ChannelReply) } return } // auth token if ok := c.AuthToken(key, token); !ok { ws.Write(AuthReply) log.Error("<%s> user_key:\"%s\" auth token \"%s\" failed", addr, key, token) return } // add a conn to the channel connElem, err := c.AddConn(key, &Connection{Conn: ws, Proto: WebsocketProto, Version: version}) if err != nil { log.Error("<%s> user_key:\"%s\" add conn error(%v)", addr, key, err) return } // blocking wait client heartbeat reply := "" begin := time.Now().UnixNano() end := begin + Second for { // more then 1 sec, reset the timer if end-begin >= Second { if err = ws.SetReadDeadline(time.Now().Add(time.Second * time.Duration(heartbeat))); err != nil { log.Error("<%s> user_key:\"%s\" websocket.SetReadDeadline() error(%v)", addr, key, err) break } begin = end } if err = websocket.Message.Receive(ws, &reply); err != nil { log.Error("<%s> user_key:\"%s\" websocket.Message.Receive() error(%v)", addr, key, err) break } if reply == Heartbeat { if _, err = ws.Write(HeartbeatReply); err != nil { log.Error("<%s> user_key:\"%s\" write heartbeat to client error(%s)", addr, key, err) break } log.Debug("<%s> user_key:\"%s\" receive heartbeat", addr, key) } else { log.Warn("<%s> user_key:\"%s\" unknown heartbeat protocol", addr, key) break } end = time.Now().UnixNano() } // remove exists conn if err := c.RemoveConn(key, connElem); err != nil { log.Error("<%s> user_key:\"%s\" remove conn error(%v)", addr, key, err) } return }
// Server Ping interface func (r *MessageRPC) Ping(p int, ret *int) error { log.Debug("ping ok") return nil }
// SubscribeTCPHandle handle the subscribers's connection. func SubscribeTCPHandle(conn net.Conn, args []string) { argLen := len(args) addr := conn.RemoteAddr().String() if argLen < 2 { conn.Write(ParamReply) log.Error("<%s> subscriber missing argument", addr) return } // key, heartbeat key := args[0] if key == "" { conn.Write(ParamReply) log.Warn("<%s> key param error", addr) return } heartbeatStr := args[1] i, err := strconv.Atoi(heartbeatStr) if err != nil { conn.Write(ParamReply) log.Error("<%s> user_key:\"%s\" heartbeat:\"%s\" argument error (%v)", addr, key, heartbeatStr, err) return } if i < minHearbeatSec { conn.Write(ParamReply) log.Warn("<%s> user_key:\"%s\" heartbeat argument error, less than %d", addr, key, minHearbeatSec) return } heartbeat := i + delayHeartbeatSec token := "" if argLen > 2 { token = args[2] } version := "" if argLen > 3 { version = args[3] } log.Info("<%s> subscribe to key = %s, heartbeat = %d, token = %s, version = %s", addr, key, heartbeat, token, version) // fetch subscriber from the channel c, err := UserChannel.Get(key, true) if err != nil { log.Warn("<%s> user_key:\"%s\" can't get a channel (%s)", addr, key, err) if err == ErrChannelKey { conn.Write(NodeReply) } else { conn.Write(ChannelReply) } return } // auth token if ok := c.AuthToken(key, token); !ok { conn.Write(AuthReply) log.Error("<%s> user_key:\"%s\" auth token \"%s\" failed", addr, key, token) return } // add a conn to the channel connElem, err := c.AddConn(key, &Connection{Conn: conn, Proto: TCPProto, Version: version}) if err != nil { log.Error("<%s> user_key:\"%s\" add conn error(%v)", addr, key, err) return } // blocking wait client heartbeat reply := []byte{0} // reply := make([]byte, HeartbeatLen) begin := time.Now().UnixNano() end := begin + Second for { // more then 1 sec, reset the timer if end-begin >= Second { if err = conn.SetReadDeadline(time.Now().Add(time.Second * time.Duration(heartbeat))); err != nil { log.Error("<%s> user_key:\"%s\" conn.SetReadDeadLine() error(%v)", addr, key, err) break } begin = end } if _, err = conn.Read(reply); err != nil { if err != io.EOF { log.Warn("<%s> user_key:\"%s\" conn.Read() failed, read heartbeat timedout error(%v)", addr, key, err) } else { // client connection close log.Warn("<%s> user_key:\"%s\" client connection close error(%v)", addr, key, err) } break } if string(reply) == Heartbeat { if _, err = conn.Write(HeartbeatReply); err != nil { log.Error("<%s> user_key:\"%s\" conn.Write() failed, write heartbeat to client error(%v)", addr, key, err) break } log.Debug("<%s> user_key:\"%s\" receive heartbeat (%s)", addr, key, reply) } else { log.Warn("<%s> user_key:\"%s\" unknown heartbeat protocol (%s)", addr, key, reply) break } end = time.Now().UnixNano() } // remove exists conn if err := c.RemoveConn(key, connElem); err != nil { log.Error("<%s> user_key:\"%s\" remove conn error(%v)", addr, key, err) } return }