From 1ac79fbc3cdad0a6dfcb44c54118fc350762fd74 Mon Sep 17 00:00:00 2001
From: dingjun <fangdingjun@gmail.com>
Date: Wed, 21 Aug 2019 15:15:27 +0800
Subject: [PATCH] fix lost response issue

---
 websocket_transport.go | 30 +++++++++++++++++++-----------
 1 file changed, 19 insertions(+), 11 deletions(-)

diff --git a/websocket_transport.go b/websocket_transport.go
index c75d75a..8dbccf9 100644
--- a/websocket_transport.go
+++ b/websocket_transport.go
@@ -100,6 +100,7 @@ func (h *WebsocketTransport) readloop() {
 			h.Mu.Unlock()
 
 			if !ok {
+				log.Warnf("handler for id %s not exists", res.ID)
 				continue
 			}
 			if res.Error != nil {
@@ -125,6 +126,7 @@ func (h *WebsocketTransport) readloop() {
 			h.Mu.Unlock()
 
 			if !ok {
+				log.Warnf("handler for method %s not exists", res.Method)
 				continue
 			}
 
@@ -172,12 +174,6 @@ func (h *WebsocketTransport) Call(method string, args interface{}, reply interfa
 		return err
 	}
 
-	log.Debugf("write %s", d)
-
-	if err := h.Conn.WriteMessage(websocket.TextMessage, d); err != nil {
-		return err
-	}
-
 	ch := make(chan json.RawMessage, 1)
 	errch := make(chan *Error, 1)
 
@@ -189,6 +185,15 @@ func (h *WebsocketTransport) Call(method string, args interface{}, reply interfa
 	}
 	h.Mu.Unlock()
 
+	log.Debugf("write %s", d)
+
+	if err := h.Conn.WriteMessage(websocket.TextMessage, d); err != nil {
+		h.Mu.Lock()
+		delete(h.inflight, id)
+		h.Mu.Unlock()
+		return err
+	}
+
 	select {
 	case data := <-ch:
 		return json.Unmarshal(data, reply)
@@ -202,11 +207,6 @@ func (h *WebsocketTransport) Call(method string, args interface{}, reply interfa
 func (h *WebsocketTransport) Subscribe(method string, notifyMethod string,
 	args interface{}, reply interface{}) (chan json.RawMessage, chan *Error, error) {
 
-	err := h.Call(method, args, reply)
-	if err != nil {
-		return nil, nil, err
-	}
-
 	ch := make(chan json.RawMessage)
 	errch := make(chan *Error)
 
@@ -218,5 +218,13 @@ func (h *WebsocketTransport) Subscribe(method string, notifyMethod string,
 	}
 	h.Mu.Unlock()
 
+	err := h.Call(method, args, reply)
+	if err != nil {
+		h.Mu.Lock()
+		delete(h.inflight, notifyMethod)
+		h.Mu.Unlock()
+		return nil, nil, err
+	}
+
 	return ch, errch, nil
 }