NSQ 源码分析之Lookupd--TCPServer

今天主要讲的是 Lookupd 中 TCP 服务的实现，TCP服务主要用于响应 NSQD 的心跳检查，NSQD 注册等信息。Lookupd 的 TCP 服务与 NSQD 中的 TCP 服务 在代码结构及实现方面非常相似。都是通过 IOLoop 循环读取客户端请求、解析命令、调用命令函数、执行命令、返回响应。

主要代码文件：

1.nsqlookupd/lookup_protocol_v1.go 

IOLoop 函数

func (p *LookupProtocolV1) IOLoop(conn net.Conn) error {
	var err error
	var line string
    //包装conn
	client := NewClientV1(conn)
	reader := bufio.NewReader(client)
	for {
        //以回车作为命令分隔
		line, err = reader.ReadString('\n')
		if err != nil {
			break
		}
        //去除左右空格
		line = strings.TrimSpace(line)
        //以空格分隔,生成命令数组
		params := strings.Split(line, " ")

		var response []byte
        //执行命令
		response, err = p.Exec(client, reader, params)
		if err != nil {
            ....
            //响应错误
			_, sendErr := protocol.SendResponse(client, []byte(err.Error()))
			if sendErr != nil {
				p.ctx.nsqlookupd.logf(LOG_ERROR, "[%s] - %s%s", client, sendErr, ctx)
				break
			}
            ...
			continue
		}

		if response != nil {
            //响应请求
			_, err = protocol.SendResponse(client, response)
			if err != nil {
				break
			}
		}
	}

	conn.Close()
	if client.peerInfo != nil {
        //清除客户信息
		registrations := p.ctx.nsqlookupd.DB.LookupRegistrations(client.peerInfo.id)
		for _, r := range registrations {
			if removed, _ := p.ctx.nsqlookupd.DB.RemoveProducer(r, client.peerInfo.id); removed {
				p.ctx.nsqlookupd.logf(LOG_INFO, "DB: client(%s) UNREGISTER category:%s key:%s subkey:%s",
					client, r.Category, r.Key, r.SubKey)
			}
		}
	}
	return err
}

 Exec 函数

func (p *LookupProtocolV1) Exec(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
	switch params[0] {
	case "PING": //心跳
		return p.PING(client, params)
	case "IDENTIFY": //NSQD 注册
		return p.IDENTIFY(client, reader, params[1:])
	case "REGISTER": //注册Topic or Channel
		return p.REGISTER(client, reader, params[1:])
	case "UNREGISTER": //注销Topic or Channel
		return p.UNREGISTER(client, reader, params[1:])
	}
	return nil, protocol.NewFatalClientErr(nil, "E_INVALID", fmt.Sprintf("invalid command %s", params[0]))
}

PING 函数 心跳检测

func (p *LookupProtocolV1) PING(client *ClientV1, params []string) ([]byte, error) {
	if client.peerInfo != nil {
		// we could get a PING before other commands on the same client connection
		cur := time.Unix(0, atomic.LoadInt64(&client.peerInfo.lastUpdate))
		now := time.Now()
        //更新心跳检测时间
		atomic.StoreInt64(&client.peerInfo.lastUpdate, now.UnixNano())
	}
	return []byte("OK"), nil
}

IDENTIFY 函数 NSQD注册

func (p *LookupProtocolV1) IDENTIFY(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
	var err error
    //防止重复注册
	if client.peerInfo != nil {
		return nil, protocol.NewFatalClientErr(err, "E_INVALID", "cannot IDENTIFY again")
	}

    ... //读取请求
    ... // 收集客户端信息，包括IP，TPC 端口，域名，广播地址等。

	client.peerInfo = &peerInfo
    //向DB注册客户端
	if p.ctx.nsqlookupd.DB.AddProducer(Registration{"client", "", ""}, &Producer{peerInfo: client.peerInfo}) {
		p.ctx.nsqlookupd.logf(LOG_INFO, "DB: client(%s) REGISTER category:%s key:%s subkey:%s", client, "client", "", "")
	}

	//建立响应，将服务相关的IP，端口等返回给客户端。
    
	return response, nil
}

REGISTER 函数注册 Topic or Channel

func (p *LookupProtocolV1) REGISTER(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
	...
    //getTopicChan 主要是解析出param 中的 topic 和 channel
	topic, channel, err := getTopicChan("REGISTER", params)
	if err != nil {
		return nil, err
	}
    //注册channel
	if channel != "" {
		key := Registration{"channel", topic, channel}
		if p.ctx.nsqlookupd.DB.AddProducer(key, &Producer{peerInfo: client.peerInfo}) {
			p.ctx.nsqlookupd.logf(LOG_INFO, "DB: client(%s) REGISTER category:%s key:%s subkey:%s",
				client, "channel", topic, channel)
		}
	}
    //注册topic
	key := Registration{"topic", topic, ""}
	if p.ctx.nsqlookupd.DB.AddProducer(key, &Producer{peerInfo: client.peerInfo}) {
		p.ctx.nsqlookupd.logf(LOG_INFO, "DB: client(%s) REGISTER category:%s key:%s subkey:%s",
			client, "topic", topic, "")
	}

	return []byte("OK"), nil
}

 UNREGISTER 函数注销 Topic or Channel

func (p *LookupProtocolV1) UNREGISTER(client *ClientV1, reader *bufio.Reader, params []string) ([]byte, error) {
    ...

	if channel != "" {
        //删除channel
		key := Registration{"channel", topic, channel}
		removed, left := p.ctx.nsqlookupd.DB.RemoveProducer(key, client.peerInfo.id)
        ...
        //删除临时节点
		if left == 0 && strings.HasSuffix(channel, "#ephemeral") {
			p.ctx.nsqlookupd.DB.RemoveRegistration(key)
		}
	} else {
        //查找 topic 下所有的channel
		registrations := p.ctx.nsqlookupd.DB.FindRegistrations("channel", topic, "*")
		for _, r := range registrations {
            //删除channel
			removed, _ := p.ctx.nsqlookupd.DB.RemoveProducer(r, client.peerInfo.id)
			if removed {
                ...
			}
		}
        //删除topic
		key := Registration{"topic", topic, ""}
		removed, left := p.ctx.nsqlookupd.DB.RemoveProducer(key, client.peerInfo.id)
        ...
        //删除临时节点
		if left == 0 && strings.HasSuffix(topic, "#ephemeral") {
			p.ctx.nsqlookupd.DB.RemoveRegistration(key)
		}
	}

	return []byte("OK"), nil
}

总结：Lookupd 中的 TCPServer 主要提供了 NSQD 注册，心跳检测、Topic or Channel 注册及 Topic or Channel 注销。

下一篇分享：Lookupd 中的 HttpServer