Chat system with Golang

邱嵩傑 included in 技術筆記

2022-02-03 835 words 4 minutes

Contents

前言

年前接了一個重構聊天系統到網頁上的任務。

現有的前端由於不是使用網頁，所以舊的 code 沒什麼參考價值。

後端的部分先前開發的 RD 已經不在公司了 orz，交接來的 RD 說花時間看舊的 code 不如寫一個新的，但他要忙別的專案沒空弄新的給我接。

Alright，後端說要寫新的，但他沒時間弄。我自己沒學過 golang，但我有的是時間，那麼身為 golang 新手，如果要弄出一個後端，邊學邊寫新的應該比直接看懂舊的來的簡單(吧?

就使用 gorilla 的 Websocket package 吧!

於是我風塵僕僕的踏進了 golang，初識了 goroutine & Channel ，並且學會一些打印出 golang 中各種資料型態的方法(相較於 js & ts，golang 這個強型態語言要從 0 開始直接理解現有的專案程式碼實在不是很容易…)。最後總算是看懂了 document 給的 example，並且延伸出了私人訊息的功能。

完成了這一切後跟老闆報告進度，老闆表示: 用舊的後端吧，現有的功能其實不只聊天，要讓我重新理解並完整重構……有點……。

well，老闆說的有理，改進比改革省時多了，而且歷經一個月的歷練，理解現有的 code 應該不再那麼難了。

但為了不要讓這曾經的努力漸漸淡忘在腦海中，還是將經歷過的紀錄下來吧!

Good to know first

channel type
make function
select for channel
Method/Function Receiver

gorilla example code

目標是將 example code 延伸出私訊的功能，所以打算先看懂現有的 code，在來開始做延伸。

main.go

首先從程式 entry point 開始看起。

func main() {
  flag.Parse()

  // create hub instance
  hub := newHub()

  // run new goroutine for hub
  go hub.run()

  // serve front end page, but it's not necessary for websocket server
  http.HandleFunc("/", serveHome)

  // websocket api
  http.HandleFunc("/ws", func(w http.ResponseWriter, r *http.Request) {
    serveWs(hub, w, r)
  })

  // start serve
  err := http.ListenAndServe(*addr, nil)
  if err != nil {
    log.Fatal("ListenAndServe: ", err)
  }
}

看起來我們所需要的 websocket 精華寫在 serveWs 以及 newHub() 中。

client.go

接著 follow 上述來到 client.go 找到 serveWs 這個 function。

func serveWs(hub *Hub, w http.ResponseWriter, r *http.Request) {
  // upgrade http request to websocket
  conn, err := upgrader.Upgrade(w, r, nil)
  if err != nil {
    log.Println(err)
    return
  }
  // declare a client variable with current connection data
  client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256)}

  // send client as a data in to register channel
  client.hub.register <- client

  // Allow collection of memory referenced by the caller by doing all work in
  // new goroutines.
  go client.writePump()
  go client.readPump()
}

資料流向的重點是 client.hub.register <- client。

並且參考官網

Channels are a typed conduit through which you can send and receive values with the channel operator, <-.

ch <- v    // Send v to channel ch.
v := <-ch  // Receive from ch, and
           // assign value to v.

從 hub.go 裡面可以知道 Hub 中的 register 是 Channel type，因此這邊可以理解是 client 這筆 data 被傳送進了 client.hub.register 這個 Channel。

hub.go

目前傳進 Channel 的 data 還沒有人接住，所以我們來到 hub.go 找到下段程式碼

// Method/Function Receiver: 'method' in golang is a 'function' with 'receiver'
// declare method 'run()' bind to '*Hub' by receiver '(h *Hub)'
func (h *Hub) run() {
  for {
    select {
    // receive message from h.register(channel type)
    case client := <-h.register:
      // update h.clients(add new client)
      h.clients[client] = true

    // client receive message from h.unregister(channel type)
    case client := <-h.unregister:
      if _, ok := h.clients[client]; ok {
      delete(h.clients, client)
      close(client.send)
    }

    // message from h.broadcast(channel type)
    case message := <-h.broadcast:
      // here the example sends message to all the clients who has already registered
      for client := range h.clients {
        select {
        // put message data in to client.send channel
        case client.send <- message:
        default:
          close(client.send)
          delete(h.clients, client)
        }
      }
    }
  }
}

注意到 case client := <-h.register: 其實就是將 register 這個 Channel 拿到的資料存進這裡的 client 變數中，然後再透過 h.clients[client] = true 來把 client 註冊進一開始在 main.go 中呼叫的 newHub() 這個 goroutine 的 clients 中，以此做到紀錄所有線上的 Clients。

如果到這邊有卡住的話可以參考上面的:

v := <-ch  // Receive from ch, and
           // assign value to v.

至此算是順利走完一趟有關 Channel 的 data flow 了。

readPump

接下來可以延續上面，看到 client.go 的最後兩行，發現 writePump 及 readPump 都各自有一個 goroutine 在處理。

go client.writePump()
go client.readPump()

首先看接收訊息的 readPump ，可以發現他在對資料做了一頓 Set 之後，最後一行執行了 c.hub.broadcast <- message，這邊同上面邏輯，可以理解是將 message 這個 data 送進 c.hub.broadcast 這個 Channel 中。

再看到 hub.go 中的 run function，case message := <-h.broadcast: 就寫出了當 h.broadcast 接收到資料時所需要做的事情，這邊發現他是將收到的資料傳送給上面記錄的 clients 裡的所有 client。

private message

第一步，紀錄訊息資料

藉由上述我們可以知道派發訊息的 feature 是透過 readPump & case message := <-h.broadcast 兩段程式碼來完成的，因此這邊可以將原本的 readPump 改寫為

func (c *Client) readPump(from string, to string) {
  defer func() {
    c.hub.unregister <- c
    c.conn.Close()
  }()
  c.conn.SetReadLimit(maxMessageSize)
  c.conn.SetReadDeadline(time.Now().Add(pongWait))
  c.conn.SetPongHandler(func(string) error {c.conn.SetReadDeadline(time.Now().Add(pongWait)); return nil })
  for {
    _, message, err := c.conn.ReadMessage()
    if err != nil {
      if websocket.IsUnexpectedCloseError(err, websocket.CloseGoingAway, websocket.CloseAbnormalClosure) {
        log.Printf("error: %v", err)
      }
      break
    }
    message = bytes.TrimSpace(bytes.Replace(message, newline, space, -1))

    // add message time
    var messageMap map[string]string
    json.Unmarshal([]byte(string(message)), &messageMap)
    nowSec := time.Now().Unix()
    messageMap["time"] = strconv.FormatInt(nowSec*1000, 10)

    // add sender and receiver
    messageMap["from"] = from
    messageMap["to"] = to
    newMessage, marshalErr := json.Marshal(messageMap)

    if marshalErr != nil {
      log.Printf("error: %v", marshalErr)
    }

    // send to hub(case message := <-h.broadcast:)
    c.hub.broadcast <- newMessage
  }
}

這邊給 readPump 增加了 from & to 兩個參數來記住 message 是來自誰以及傳給誰，並且同時增加了 time field 來記錄伺服器傳送訊息的時間。

第二步，連接 ws 時紀錄 sender & receiver

在呼叫 readPump 的 serveWs 中增加接收"傳送訊息人( sender )“以及"接收訊息人( receiver )“的程式碼，如下。

senderId := r.URL.Query().Get("senderId")
receiverId := r.URL.Query().Get("receiverId")
client := &Client{hub: hub, conn: conn, send: make(chan []byte, 256), senderId: senderId}
// ---------------
go client.readPump(senderId, messageReceiver)

由於上面在 Client type 中新增了 senderId 這個 field，因此宣告 Client structure 的地方也要改為

type Client struct {
	hub *Hub

	// The websocket connection.
	conn *websocket.Conn

	// Buffered channel of outbound messages.
	send chan []byte

	// add senderId field
	senderId string
}

第三步，只把訊息傳給指定的 sender & receiver

接著來到 hub.go 裡的 case message := <-h.broadcast:

case message := <-h.broadcast:
  for client := range h.clients {
    // this part looks weird, may need golang pro to improve
    var messageMap map[string]string
    json.Unmarshal([]byte(string(message)), &messageMap)

    // find specific sender and receiver in clients map
    if messageMap["to"] == client.senderId || messageMap["from"] == client.senderId {
      select {
      // receive message from socket
        case client.send <- message:
        default:
          close(client.send)
          delete(h.clients, client)
      }
    }
  }
}