golang使用channel实现broadcast

        最近的一个项目需要一个类似于广播的功能, 当客户端连接到服务端以后, 需要把相同的一份数据分别发送给每个客户端, 功能虽然简单, 但是还是有坑, 主要的坑就是channel阻塞的问题. 所本文模拟了这个服务, 加深自己的理解. 

        代码运行环境为: Linux 3.10.0-957.21.3.el7.x86_64

        代码结构如下:

        其中broad.go文件, 功能主要是监听client的注册(RegisterChan), 注销(UnregisterChan), 以及定时(time.NewTicker)这三个事件并维护一个map, 代码如下:

package broadcast

import (
	"time"
)

//Broad ..
type Broad struct {
	ClientsMap     map[*Client]bool // map
	RegisterChan   chan *Client     // 连接注册
	UnregisterChan chan *Client     // 关闭连接
}

// NewBroad ..
func NewBroad() *Broad {
	return &Broad{
		ClientsMap:     make(map[*Client]bool),
		RegisterChan:   make(chan *Client, 1),
		UnregisterChan: make(chan *Client, 1),
	}
}

// Run ..
func (broad *Broad) Run() {
	defer func() {
		close(broad.RegisterChan)
		close(broad.UnregisterChan)
	}()

	ticker := time.NewTicker(time.Microsecond * 1)
	defer ticker.Stop()

	msg := []byte("HI.")

	for {
		select {
		case client := <-broad.RegisterChan:
			broad.ClientsMap[client] = true

		case client := <-broad.UnregisterChan:
			client.Conn.Close()              //关闭连接
			delete(broad.ClientsMap, client) // 删除元素
			close(client.WriteChan)          // 关闭channel

		case <-ticker.C:
			for client := range broad.ClientsMap {
				client.WriteChan <- &msg
			}
		}
	}
}

        client.go文件, 主要是接收WriteChan的事件, 并往客户端写数据且设置了超时时间, 其中为了模拟我遇到的坑, 增加了一个userGo的选项, 可以自行模拟, 代码如下:

package broadcast

import (
	"log"
	"net"
	"time"
)

// Client ..
type Client struct {
	Broad     *Broad
	Conn      net.Conn
	WriteChan chan *[]byte
}

// WriteMsg ..
func (client *Client) WriteMsg() {

	defer func() {
		// 代码运行到这里的时候, 还没有向Broad协程发送注销channel,
		// ClientsMap里还是有该client. 当写该WriteChan时, 会造成
		// Broad协程阻塞, 无法再监听新的连接. 此时需要开一个新的协
		// 程继续监听WriteChan事件, 直到Broad协程关闭了WriteChan.

		// 可以不使用这段代码, 并把Broad协程里的秒改成毫秒/微秒并断
		// 掉客户端的网络进行尝试.
		if *useGo {
			go func() {
				for {
					select {
					case msg, ok := <-client.WriteChan:
						if !ok || msg == nil {
							log.Println("Is not ok.")
							return
						}
					}
				}
			}()
		}

		// 通知Broad, 此client离线
		client.Broad.UnregisterChan <- client
	}()

	for {
		select {
		case msg, ok := <-client.WriteChan:
			if !ok || msg == nil {
				log.Println("Is not ok or equal nil.")
				return
			}

			if err := client.Conn.SetWriteDeadline(time.Now().Add(time.Second)); err != nil {
				log.Println("SetWriteDeadline failed. ErrMsg: " + err.Error())
				return
			}
			if _, err := client.Conn.Write(*msg); err != nil {
				log.Println("Write failed. ErrMsg: " + err.Error())
				return
			}
		}
	}
}

        server.go文件, 功能主要是实现TCP服务, 代码如下:

package broadcast

import (
	"flag"
	"log"
	"net"
)

var useGo = flag.Bool("use_go", false, "是否使用新协程读取WriteChan, 防止阻塞")
var addr = flag.String("addr", "localhost:8080", "TCP addr")

func init() {
	flag.Parse()
}

// Start ..
func Start() {
	broad := NewBroad()
	go broad.Run()

	listener, err := net.Listen("tcp", *addr)
	if err != nil {
		log.Fatal("Listen fail.")
	}
	log.Println("Listen success.")
	defer listener.Close()

	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Println("Accept failed. ErrMsg: " + err.Error())
			continue
		}
		log.Println("Accept successfully. Remote addr: " + conn.RemoteAddr().String())

		client := &Client{Broad: broad, Conn: conn, WriteChan: make(chan *[]byte)}
		client.Broad.RegisterChan <- client

		go client.WriteMsg()
	}
}

        main.go文件主要是调用这个包, 这里不再展示.

        模拟遇到的坑结果如下:

图片左边的客户端的网络一直都是好的, 而在我断掉了第二个客户端的网络后, Broad协程阻塞在client.WriteChan <- &msg这, 导致第一个客户端的也无法收到数据. 解决办法就是重新开启一个协程接收WriteChan, 直到Broad协程接收到UnregisterChan事件并关闭客户端.