Golangでchannelを使いたい - Adwaysエンジニアブログ

Adways Advent Calendar 1日目の記事です。

http://blog.engineer.adways.net/entry/advent_calendar/archive

Adwaysアドベントカレンダー１日目を担当します、安藤です。
業務では主にrubyを使っていますが、今回はchannelを使いたかったのでgoにしました。(笑)

今回学習として簡素なチャットを作ってみました。
channelを使った箇所は、

入退室管理
メッセージ送信

です。

モデル

今回は一つのroomで複数のクライアントが接続するのを想定してます。（最大でどれほどのクライアントが作れるかは試してません。）

package main                                                                                                                                                                         

import (
    "github.com/gorilla/websocket"
    "log"
    "net/http"
)

// クライアント
type client struct {
    socket *websocket.Conn   // WebSocketのコネクション
    send   chan []byte       // メッセージをブラウザに送信するchannel
    room   *room             // クライアントが属するチャットルーム
}

// チャットルーム
type room struct {
    forward chan []byte      // 他のすべてのクライアントに送信するメッセージを持つ
    join    chan *client     // 入室するクライアント
    leave   chan *client     // 退室するクライアント
    clients map[*client]bool // 入室中のすべてのクライアント
}

// チャットルームの初期化処理
func newRoom() *room {
    return &room{
        forward: make(chan []byte),
        join:    make(chan *client),
        leave:   make(chan *client),
        clients: make(map[*client]bool),
    }   
}

client
* send ・・・・ここにはブラウザに送信するメッセージをキューのように保持します。
room
* join, leave, clients ・・・・ clientsをそのままいじるのではなく、追加・削除もchannelを通して行います。

ブラウザへの送受信部分

read
* ブラウザからのメーセージをWebSocket経由で受け取り、foward channelに渡しています。
つまり、自分が書いたメッセージを他のユーザーに伝えるために、メッセージを蓄えさせているところ。
write
* send channelに届いたメッセージをブラウザに送信しています。

func (c *client) read() {
    for {
        if _, msg, err := c.socket.ReadMessage(); err == nil {
            c.room.forward <- msg
        } else {
            break
        }
    }
    c.socket.Close()
}

func (c *client) write() {
    for msg := range c.send {
        if err := c.socket.WriteMessage(websocket.TextMessage, msg); err != nil {
            break
        }
    }
    c.socket.Close()
}

並行処理

runメソッドは無限ループさせており、強制終了するまでjoin, leave, forawardを監視します。
どれかのchannelにメッセージが届くと、そのcase文が実行されるしくみです。

func (r *room) run() {
    for {
        select {
        case client := <-r.join:
            //入室
            r.clients[client] = true
        case client := <-r.leave:
            delete(r.clients, client)
            close(client.send)
        case msg := <-r.forward:
            // すべてのクライアント(ユーザー)にメッセージを送信する
            for client := range r.clients {
                select {
                case client.send <- msg:
                    // 送信
                default:
                    // 失敗
                    delete(r.clients, client)
                    close(client.send)
                }
            }
        }
    }
}

select caseを使うと同時に実行される危険がないので、clientsへの変更が同時に発生することがないです。

defaultでは、send channelに送信できなかった場合、クライアントの退室処理とsend channelをcloseしてます。

ちなみに、closeしたchannelへの送信はランタイムパニックを発生させます。
(closeしたchannelにはアクセスしないようにしましょう。)

その他

roomをHTTPハンドラにし、ブラウザにアクセス時クライアントを生成、
client.read()をメインスレッドで行うことでコネクションを保持し、クライアントの終了時、退室処理をするように追加しました。

const (
    socketBufferSize  = 1024
    messageBufferSize = 256
)

// WebSocketを利用するなら必要
var upgrader = &websocket.Upgrader{ReadBufferSize: socketBufferSize, WriteBufferSize: socketBufferSize}

func (r *room) ServeHTTP(w http.ResponseWriter, req *http.Request) {
    socket, err := upgrader.Upgrade(w, req, nil)
    if err != nil {
        log.Fatal("ServeHTTP:", err)
        return
    }
    client := &client{
        socket: socket,
        send:   make(chan []byte, messageBufferSize), // バッファサイズを設定
        room:   r,
    }
    r.join <- client
    defer func() { r.leave <- client }() // この関数が終了する時に呼ばれる
    go client.write()                    // 別のスレッドで書き出す
    client.read()                        // メインスレッドで読み込み
}