package main
import (
  "fmt"
  "net/http"
)
func main() {
 fmt.Println("Listen localhost:8080")
   // 注冊(cè)一個(gè)用于WebSocket的路由，實(shí)際業(yè)務(wù)中不可能只有一個(gè)路由
  http.HandleFunc("/messages", messageHandler)
  // 監(jiān)聽8080端口，沒有實(shí)現(xiàn)服務(wù)異常處理器，因此第二個(gè)參數(shù)是nil
  http.ListenAndServe("localhost:8080", nil)
}
func messageHandler(response http.ResponseWriter, request *http.Request) {
  // TODO: 實(shí)現(xiàn)消息處理
  response.Write([]byte("HelloWorld"))
}

func messageHandler(response http.ResponseWriter, request *http.Request) {
  var upgrader = websocket.Upgrader{
    // 允許跨域
    CheckOrigin: func(resquest *http.Request) bool {
      return true
    },
  }
  // 建立連接
  conn, err := upgrader.Upgrade(response, request, nil)
  if err != nil {
    return
  }
  // 收發(fā)消息
  for {
    // 讀取消息
    _, bytes, err := conn.ReadMessage()
    if err != nil {
      _ = conn.Close()
    }
    // 寫入消息
    err = conn.WriteMessage(websocket.TextMessage, bytes)
    if err != nil {
      _ = conn.Close()
    }
  }
}

// main.go
package main
import (
  "bluemiaomiao.github.io/websocket-go/service"
  "fmt"
  "net/http"
  "github.com/gorilla/websocket"
)
func main() {
  fmt.Println("Listen localhost:8080")
  http.HandleFunc("/messages", messageHandler)
  _ = http.ListenAndServe("localhost:8080", nil)
}
func messageHandler(response http.ResponseWriter, request *http.Request) {
  var upgrader = websocket.Upgrader{
    // 允許跨域
    CheckOrigin: func(resquest *http.Request) bool {
      return true
    },
  }
  // 建立連接
  conn, err := upgrader.Upgrade(response, request, nil)
  wsConn, err := service.Create(conn)
  if err != nil {
    return
  }
  // 收發(fā)消息
  for {
    // 讀取消息
    msg, err := wsConn.ReadOne()
    if err != nil {
      wsConn.Close()
    }
    // 寫入消息
    err = wsConn.WriteOne(msg)
    if err != nil {
      _ = conn.Close()
    }
  }
}

// service/messsage_service.go
package service
import (
  "errors"
  "github.com/gorilla/websocket"
  "sync"
)
// 封裝的連接對(duì)象
// 
// 由于websocket的Close()方法是可重入的，所以可以多次調(diào)用，但是關(guān)閉Channel的close()
// 方法不是可重入的，因此通過isClosed進(jìn)行判斷
// isClosed可能發(fā)生資源競(jìng)爭(zhēng)，因此通過互斥鎖避免
// 關(guān)閉websocket連接后，也要自動(dòng)關(guān)閉輸入輸出消息流，因此通過signalCloseLoopChan實(shí)現(xiàn)
type Connection struct {
  conn                  *websocket.Conn  // 具體的連接對(duì)象
  inputStream             chan []byte       // 輸入流，使用Channel模擬
  outputStream           chan []byte       // 輸出流，使用chaneel模擬
  signalCloseLoopChan     chan byte       // 關(guān)閉信號(hào)
  isClosed               bool            // 是否調(diào)用過close()方法
  lock                   sync.Mutex      // 簡(jiǎn)單的鎖
}
// 用于初始化一個(gè)連接對(duì)象
func Create(conn *websocket.Conn) (connection *Connection, err error) {
  connection = &Connection{
    conn:              conn,
    inputStream:        make(chan []byte, 1000),
    outputStream:       make(chan []byte, 1000),
    signalCloseLoopChan: make(chan byte, 1),
    isClosed:            false,
  }
  // 啟動(dòng)讀寫循環(huán)
  go connection.readLoop()
  go connection.writeLoop()
  return
}
// 讀取一條消息
func (c *Connection) ReadOne() (msg []byte, err error) {
  select {
  case msg = <-(*c).inputStream:
  case <-(*c).signalCloseLoopChan:
    err = errors.New("connection is closed")
  }
  return
}
// 寫入一條消息
func (c *Connection) WriteOne(msg []byte) (err error) {
  select {
  case (*c).outputStream <- msg:
  case <-(*c).signalCloseLoopChan:
    err = errors.New("connection is closed")
  }
  return
}
// 關(guān)閉連接對(duì)象
func (c *Connection) Close() {
  _ = (*c).conn.Close()
  (*c).lock.Lock()
  if !(*c).isClosed {
    close((*c).signalCloseLoopChan)
  }
  (*c).lock.Unlock()
}
// 讀取循環(huán)
func (c *Connection) readLoop() {
  // 不停的讀取長(zhǎng)連接中的消息，只要存在消息就將其放到隊(duì)列中
  for {
    _, bytes, err := (*c).conn.ReadMessage()
    if err != nil {
      (*c).Close()
    }
    select {
    case <-(*c).signalCloseLoopChan:
      (*c).Close()
    case (*c).inputStream <- bytes:
    }
  }
}
// 寫入循環(huán)
func (c *Connection) writeLoop() {
  // 只要隊(duì)列中存在消息，就將其寫入
  var data []byte
  for {
    select {
    case data = <-(*c).outputStream:
    case <-(*c).signalCloseLoopChan:
      (*c).Close()
    }
    err := (*c).conn.WriteMessage(websocket.TextMessage, data)
    if err != nil {
      _ = (*c).conn.Close()
    }
  }
}