go语言小练习——基于goroutine实现的Tcp聊天室

前言

博主最近没怎么写go，最近正好放暑假，写了一个小demo来复习一下，源码会放在资源了，大家按需取用。

服务端

package main

import (
	"bufio"
	"fmt"
	"github.com/sirupsen/logrus"
	"net"
	"strconv"
	"time"
)

type User struct {
	UID         int
	Addr        string
	EnterAt     time.Time
	MessageChan chan Message
}

type Message struct {
	Owner   int    //通过uid来锁定用户
	Content string // 消息内容
}

const MaxNum = 100 //最大用户数
var UserIDMap = make([]int, MaxNum)
var (
	EnterChan   = make(chan *User)      //登记用户
	LeaveChan   = make(chan *User)      //删除用户
	messageChan = make(chan Message, 8) //广播消息
)

func GenUserID() int {
	for i := 0; i < MaxNum; i++ {
		if UserIDMap[i] == 0 {
			UserIDMap[i] = 1
			return i
		}
	}
	return -1
}

// 记录聊天室用户，并且进行广播
func broadcaster() {
	fmt.Println("聊天室启动")
	var users = make([]*User, MaxNum)
	for {
		select {
		case user := <-EnterChan:
			users[user.UID] = user
		case user := <-LeaveChan:
			users[user.UID] = nil
			UserIDMap[user.UID] = 0
			close(user.MessageChan)
		case message := <-messageChan:
			for _, user := range users {
				if user != nil && user.UID == message.Owner {
					user.MessageChan <- message
				}
			}
		}
	}
}

func handleConn(coon net.Conn) {
	defer coon.Close()
	id := GenUserID()
	if id == -1 {
		logrus.Warning("当前用户连接数超过最大限制")
		return
	}
	//创建用户实例
	user := &User{
		UID:         id,
		Addr:        coon.RemoteAddr().String(),
		EnterAt:     time.Now(),
		MessageChan: make(chan Message, 100),
	}
	// 开一个协程用来接收消息
	go SendMessage(coon, user.MessageChan)

	//向该用户发送欢迎消息,向全体用户广播消息
	messageChan <- Message{
		Owner:   user.UID,
		Content: fmt.Sprintf("%s进入聊天室", strconv.Itoa(user.UID)),
	}
	user.MessageChan <- Message{
		Owner:   user.UID,
		Content: fmt.Sprintf("%s进入聊天室", strconv.Itoa(user.UID)),
	}

	//将该用户加入到用户列表中
	EnterChan <- user

	//之前我们用另一个协程来完成有关于写(发送)的部分，这部分我们所要完成的就是有关于读取(接收)的部分，
	input := bufio.NewScanner(coon)
	for input.Scan() {
		messageChan <- Message{
			Owner:   user.UID,
			Content: fmt.Sprintf("%s:%s", strconv.Itoa(user.UID), input.Text()),
		}
	}
	if err := input.Err(); err != nil {
		logrus.Error("读取错误", err)
	}
	//用户离开聊天室
	LeaveChan <- user
	messageChan <- Message{
		Owner:   user.UID,
		Content: fmt.Sprintf("%s离开聊天室", strconv.Itoa(user.UID)),
	}
}

func SendMessage(coon net.Conn, messageChan <-chan Message) {
	for mes := range messageChan {
		_, _ = fmt.Fprintln(coon, mes)
	}
}

func main() {
	listen, err := net.Listen("tcp", ":2020")
	if err != nil {
		panic(err)
	}
	go broadcaster()
	for {
		coon, err := listen.Accept()
		if err != nil {
			logrus.Error(err)
			continue
		}
		go handleConn(coon)
	}
}

客户端

package main

import (
	"github.com/sirupsen/logrus"
	"io"
	"net"
	"os"
	"sync"
)

func main() {
	coon, err := net.Dial("tcp", "localhost:2020")
	if err != nil {
		logrus.Fatalf("Failed to connect to the server: %v", err)
	}
	defer coon.Close()

	// 用于同步两个io.Copy操作的完成
	var wg sync.WaitGroup

	// 从服务器接收消息并写入标准输出
	wg.Add(1)
	go func() {
		defer wg.Done()
		if _, err := io.Copy(os.Stdout, coon); err != nil && !isEOF(err) {
			logrus.Errorf("Failed to read from server: %v", err)
		}
	}()

	// 从标准输入读取数据并写入服务器
	wg.Add(1)
	go func() {
		defer wg.Done()
		if _, err := io.Copy(coon, os.Stdin); err != nil && !isEOF(err) {
			logrus.Errorf("Failed to write to server: %v", err)
		}
	}()

	// 等待两个io.Copy操作完成
	wg.Wait()
}

// 检查错误是否是io.EOF，这通常表示流的正常结束
func isEOF(err error) bool {
	return err == io.EOF
}

由于我是windows系统下写的，编译脚本用的就是.bat了，编译脚本如下：

REM 清理旧的编译结果
 del *.exe

REM 编译Go程序
 go build -o server.exe server.go
 go build -o client.exe client.go

然后在指定目录下运行.exe就可以了

 ./client.exe
 ./server.exe

拓展

大家如果觉得还不够的话，博主有两个功能懒得写了，大家可以试试：

• 将用户分为活跃用户和潜水用户
• 将潜水用户自动剔除群聊
  好了，大家可以试试哦,拜拜！