基於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 通訊；