基于TCP实现简单的聊天室

原文出处：《Go 语言编程之旅》第四章4.1节

基于TCP的聊天室

1、服务端

• 新用户到来，生成一个User的实例，代表该用户。

type User struct{
	ID int   // 用户的唯一标识，通过GenUserID 函数生成
	Addr string   // 用户的IP地址和端口
	EnterAt time.Time   // 用户进入的时间
	MessageChannel chan string  // 当前用户发送消息的通道
}

• 新开一个goroutine用于给用户发送消息

func sendMessage(conn net.Conn, ch <- chan string){
	for msg := range ch{
		fmt.Fprintln(conn, msg)
	}
}

结合User结构体的MessageChannel,很容易知道，需要给某个用户发送消息，只需要往该用户的MessageChannel中写入消息即可。这里需要特别提醒下，因为sendMessage在一个新的goroutine中，如果函数的ch不关闭，该goroutine是不会退出的，因此需要注意不关闭ch导致goroutine泄露问题。

• 给当前用户发送欢迎信息，同时给聊天室所有的用户发送有新用户到来的提醒

user.MessageChannel <- "Welcome" + user.String()
msg := Message{
    OwnerID: user.ID,
    Content: "user:`" + strconv.Itoa(user.ID) + "` has enter",
}
messageChannel <- msg

• 将该新用户写入全局用户列表，也就是聊天室用户列表。同时控制用户超时退出，超过5分钟没有任何响应，则提出

enteringChannel <- user

// 控制超时用户踢出
var userActive = make(chan struct{})
go func() {
    d := 5 * time.Minute
    timer := time.NewTimer(d)
    for{
        select {
            case <- timer.C:
            conn.Close()
            case <- userActive:
            timer.Reset(d)
        }
    }
}()

• 读取用户的输入，并将用户信息发送给其他用户。

  在bufio包中有多重方式获取文本输入，ReadBytes、ReadString和独特的ReadLine,对于简单的目的这些都有些复杂。在Go1,1中添加了一个新类型，Scabber,以便更容易的处理如按行读取输入序列或空格分隔单词等这类简单任务。它终结了如输入一个很长的有问题的行这样的输入错误，并且提供了简单的默认行为：基于行的输入，每行都提出了分隔标识。

//  循环读取用户的输入
input := bufio.NewScanner(conn)
for input.Scan(){
    msg.Content = strconv.Itoa(user.ID) + ";" + input.Text()
    messageChannel <- msg

    // 用户活跃
    userActive <- struct{}{}
}

if err := input.Err();err != nil {
    log.Println("读取错误：", err)
}

• 用户离开，需要做登记，并给连天使其他用户发通知

leavingChannel <- user
msg.Content =  "user: `" + strconv.Itoa(user.ID) + "` has left"
messageChannel <- msg

完整代码

package main

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"strconv"
	"sync"
	"time"
)

type User struct{
	ID int   // 用户的唯一标识，通过GenUserID 函数生成
	Addr string   // 用户的IP地址和端口
	EnterAt time.Time   // 用户进入的时间
	MessageChannel chan string  // 当前用户发送消息的通道
}

// 给用户发送信息
type Message struct{
	OwnerID int
	Content string
}

var (
	// 新用户到来，通过该channel进行登记
	enteringChannel = make(chan *User)
	// 用户离开，通过该channel进行登记
	leavingChannel = make(chan *User)
	// 广播专用的用户普通消息channel, 缓冲是尽可能避免出现异常情况阻塞
	messageChannel = make(chan Message, 9)
)

func (u *User) String() string{
	return u.Addr + ",UID:" + strconv.Itoa(u.ID) + ", Enter At:" + u.EnterAt.Format("2006-01-02 15:04:05+8000")
}

func main() {
	listener, err := net.Listen("tcp",":2020")
	if err != nil {
		panic(err)
	}

	go broadcaster()

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println(err)
			continue
		}

		go handleConn(conn)
	}
}

// broadcaster 用于记录聊天室用户，并进行消息广播
// 1. 新用户进来; 2.用户普通消息; 3.用户离开
func broadcaster(){
	users := make(map[*User]struct{})

	for {
		select{
		case user := <- enteringChannel:
			// 新用户进入
			users[user] = struct{}{}
		case user := <- leavingChannel:
			// 用户离开
			delete(users, user)
			// 避免goroutine泄露
			close(user.MessageChannel)
		case msg := <-messageChannel:
			// 给所有在线用户发送消息
			for user := range users {
				if user.ID == msg.OwnerID{
					continue
				}
				user.MessageChannel <- msg.Content
			}
		}
	}
}

func handleConn(conn net.Conn){
	defer conn.Close()

	// 1. 新用户进来，构建该用户实例
	user := &User{
		ID: GenUserID(),
		Addr: conn.RemoteAddr().String(),
		EnterAt: time.Now(),
		MessageChannel: make(chan string,8),
	}

	// 2. 当前在一个新的goroutine 中，用来进行读写操作，因此需要开一个goroutine用于读写操作
	// 读写goroutine 之间通过channel 进行通信
	go sendMessage(conn, user.MessageChannel)

	// 3. 给当前用户发送欢迎信息；给所有用户告知新用户列表
	user.MessageChannel <- "Welcome" + user.String()
	msg := Message{
		OwnerID: user.ID,
		Content: "user:`" + strconv.Itoa(user.ID) + "` has enter",
	}
	messageChannel <- msg

	// 4. 将该记录到全局的用户列表中，避免用锁
	enteringChannel <- user

	// 控制超时用户踢出
	var userActive = make(chan struct{})
	go func() {
		d := 5 * time.Minute
		timer := time.NewTimer(d)
		for{
			select {
			case <- timer.C:
				conn.Close()
			case <- userActive:
				timer.Reset(d)
			}
		}
	}()

	// 5. 循环读取用户的输入
	input := bufio.NewScanner(conn)
	for input.Scan(){
		msg.Content = strconv.Itoa(user.ID) + ";" + input.Text()
		messageChannel <- msg

		// 用户活跃
		userActive <- struct{}{}
	}

	if err := input.Err();err != nil {
		log.Println("读取错误：", err)
	}

	// 6. 用户离开
	leavingChannel <- user
	msg.Content =  "user: `" + strconv.Itoa(user.ID) + "` has left"
	messageChannel <- msg
}

func sendMessage(conn net.Conn, ch <- chan string){
	for msg := range ch{
		fmt.Fprintln(conn, msg)
	}
}

// 生成用户id
var (
	globalID int
	idocker sync.Mutex
)

func GenUserID() int {
	idocker.Lock()
	defer idocker.Unlock()

	globalID ++
	return globalID
}

2、客户端

客户端的实现直接采用 《The Go Programming Language》一书对应的示例源码：ch8/netcat3/netcat.go 。

func main() {
	conn, err := net.Dial("tcp", ":2020")
	if err != nil {
		panic(err)
	}

	done := make(chan struct{})
	go func() {
		io.Copy(os.Stdout, conn) // NOTE: ignoring errors
		log.Println("done")
		done <- struct{}{} // signal the main goroutine
	}()

	mustCopy(conn, os.Stdin)
	conn.Close()
	<-done
}

func mustCopy(dst io.Writer, src io.Reader) {
	if _, err := io.Copy(dst, src); err != nil {
		log.Fatal(err)
	}
}

• 新开了一个 goroutine 用于接收消息；
• 通过 io.Copy 来操作 IO，包括从标准输入读取数据写入 TCP 连接中，以及从 TCP 连接中读取数据写入标准输出；
• 新开的 goroutine 通过一个 channel 来和 main goroutine 通讯；