package main
import (
  "fmt"
  "net/http"
)
func main() {
 fmt.Println("Listen localhost:8080")
   // 注冊一個用于WebSocket的路由，實際業(yè)務中不可能只有一個路由
  http.HandleFunc("/messages", messageHandler)
  // 監(jiān)聽8080端口，沒有實現(xiàn)服務異常處理器，因此第二個參數(shù)是nil
  http.ListenAndServe("localhost:8080", nil)
}
func messageHandler(response http.ResponseWriter, request *http.Request) {
  // TODO: 實現(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"
)
// 封裝的連接對象
// 
// 由于websocket的Close()方法是可重入的，所以可以多次調用，但是關閉Channel的close()
// 方法不是可重入的，因此通過isClosed進行判斷
// isClosed可能發(fā)生資源競爭，因此通過互斥鎖避免
// 關閉websocket連接后，也要自動關閉輸入輸出消息流，因此通過signalCloseLoopChan實現(xiàn)
type Connection struct {
  conn                  *websocket.Conn  // 具體的連接對象
  inputStream             chan []byte       // 輸入流，使用Channel模擬
  outputStream           chan []byte       // 輸出流，使用chaneel模擬
  signalCloseLoopChan     chan byte       // 關閉信號
  isClosed               bool            // 是否調用過close()方法
  lock                   sync.Mutex      // 簡單的鎖
}
// 用于初始化一個連接對象
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,
  }
  // 啟動讀寫循環(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
}
// 關閉連接對象
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() {
  // 不停的讀取長連接中的消息，只要存在消息就將其放到隊列中
  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() {
  // 只要隊列中存在消息，就將其寫入
  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()
    }
  }
}