From 827f8e6797dd322ffd5d8e132740c72bccf5623c Mon Sep 17 00:00:00 2001
From: clchen <clchen@leyantech.com>
Date: Tue, 19 Jul 2016 12:12:13 +0800
Subject: [PATCH] fix goroutine safe

---
 wetsock/wetsock.go | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/wetsock/wetsock.go b/wetsock/wetsock.go
index 84b998b..6374384 100644
--- a/wetsock/wetsock.go
+++ b/wetsock/wetsock.go
@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"errors"
 	"reflect"
+	"sync"
 
 	"github.com/gorilla/websocket"
 	"github.com/tv42/birpc"
@@ -11,6 +12,10 @@ import (
 
 type codec struct {
 	WS *websocket.Conn
+	// https://godoc.org/github.com/gorilla/websocket#hdr-Concurrency
+	// As above document.Only one concurrent reader and one concurrent writer are allowed.
+	readMu sync.Mutex
+	writeMu sync.Mutex
 }
 
 // This is ugly, but i need to override the unmarshaling logic for
@@ -27,6 +32,9 @@ type jsonMessage struct {
 }
 
 func (c *codec) ReadMessage(msg *birpc.Message) error {
+	c.readMu.Lock()
+	defer c.readMu.Unlock()
+
 	var jm jsonMessage
 	err := c.WS.ReadJSON(&jm)
 	if err != nil {
@@ -41,6 +49,9 @@ func (c *codec) ReadMessage(msg *birpc.Message) error {
 }
 
 func (c *codec) WriteMessage(msg *birpc.Message) error {
+	c.writeMu.Lock()
+	defer c.writeMu.Unlock()
+
 	return c.WS.WriteJSON(msg)
 }