Merge pull request #88 from DaniloMoura1/refactor-websocket

Refactor websocker.go & websocket_server.go to merge both goroutines
jcelliott · Sep 2, 2015 · 3fb89b0 · 3fb89b0
2 parents cd69737 + 630cfdd
commit 3fb89b0
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 45 deletions.
diff --git a/websocket.go b/websocket.go
@@ -45,23 +45,8 @@ func newWebsocketPeer(url, protocol, origin string, serializer Serializer, paylo
 		serializer:  serializer,
 		payloadType: payloadType,
 	}
-	go func() {
-		for {
-			// TODO: use conn.NextMessage() and stream
-			// TODO: do something different based on binary/text frames
-			if _, b, err := conn.ReadMessage(); err != nil {
-				conn.Close()
-				break
-			} else {
-				msg, err := serializer.Deserialize(b)
-				if err != nil {
-					// TODO: handle error
-				} else {
-					ep.messages <- msg
-				}
-			}
-		}
-	}()
+	go ep.run()
+
 	return ep, nil
 }
 
@@ -85,3 +70,32 @@ func (ep *websocketPeer) Close() error {
 	ep.closed = true
 	return ep.conn.Close()
 }
+
+func (ep *websocketPeer) run() {
+	for {
+		// TODO: use conn.NextMessage() and stream
+		// TODO: do something different based on binary/text frames
+		if msgType, b, err := ep.conn.ReadMessage(); err != nil {
+			if ep.closed {
+				log.Println("peer connection closed")
+			} else {
+				log.Println("error reading from peer:", err)
+				ep.conn.Close()
+			}
+			close(ep.messages)
+			break
+		} else if msgType == websocket.CloseMessage {
+			ep.conn.Close()
+			close(ep.messages)
+			break
+		} else {
+			msg, err := ep.serializer.Deserialize(b)
+			if err != nil {
+				log.Println("error deserializing peer message:", err)
+				// TODO: handle error
+			} else {
+				ep.messages <- msg
+			}
+		}
+	}
+}
diff --git a/websocket_server.go b/websocket_server.go
@@ -140,33 +140,7 @@ func (s *WebsocketServer) handleWebsocket(conn *websocket.Conn) {
 		messages:    make(chan Message, 10),
 		payloadType: payloadType,
 	}
-	go func() {
-		for {
-			// TODO: use conn.NextMessage() and stream
-			// TODO: do something different based on binary/text frames
-			if msgType, b, err := conn.ReadMessage(); err != nil {
-				if peer.closed {
-					log.Println("peer connection closed")
-				} else {
-					log.Println("error reading from peer:", err)
-					conn.Close()
-				}
-				close(peer.messages)
-				break
-			} else if msgType == websocket.CloseMessage {
-				conn.Close()
-				close(peer.messages)
-				break
-			} else {
-				msg, err := serializer.Deserialize(b)
-				if err != nil {
-					log.Println("error deserializing peer message:", err)
-					// TODO: handle error
-				} else {
-					peer.messages <- msg
-				}
-			}
-		}
-	}()
+	go peer.run()
+
 	s.Router.Accept(&peer)
 }