ing

central/client/client.go

@@ -54,31 +54,30 @@ type Client interface {
 	Dial(url string, header http.Header, readBufSize int, writeBufSize int) (*http.Response, error)
 	Call(method string, args interface{}, result interface{}) error
 	Notify(method string, args interface{}) error
-	OnNotify(method string, cb OnNotifyFunc)
+	OnNotify(cb OnNotifyFunc)
 	OnClose(cb OnCloseFunc)
 	Shutdown(ctx context.Context) error
 }
 
 type client struct {
-	conn             *websocket.Conn
-	sendMutex        sync.Mutex
-	request          protocol.Request
-	notification     protocol.Notification
-	mutex            sync.Mutex
-	requestID        uint64
-	pending          map[uint64]*Call
-	closing          bool // user has called Close
-	shutdown         bool // server has told us to stop
-	onNotifyHandlers map[string][]OnNotifyFunc
-	onCloseHandlers  []OnCloseFunc
+	conn            *websocket.Conn
+	sendMutex       sync.Mutex
+	request         protocol.Request
+	notification    protocol.Notification
+	mutex           sync.Mutex
+	requestID       uint64
+	pending         map[uint64]*Call
+	closing         bool // user has called Close
+	shutdown        bool // server has told us to stop
+	onNotifyHandler OnNotifyFunc
+	onCloseHandlers []OnCloseFunc
 }
 
 func New() Client {
 	c := &client{
-		requestID:        0,
-		pending:          make(map[uint64]*Call),
-		onNotifyHandlers: make(map[string][]OnNotifyFunc),
-		onCloseHandlers:  make([]OnCloseFunc, 1),
+		requestID:       0,
+		pending:         make(map[uint64]*Call),
+		onCloseHandlers: make([]OnCloseFunc, 1),
 	}
 
 	return c
@@ -123,15 +122,8 @@ func (c *client) Notify(method string, args interface{}) error {
 	return nil
 }
 
-func (c *client) OnNotify(method string, cb OnNotifyFunc) {
-	var hs []OnNotifyFunc
-	var ok bool
-	if hs, ok = c.onNotifyHandlers[method]; !ok {
-		hs = make([]OnNotifyFunc, 1)
-		c.onNotifyHandlers[method] = hs
-	}
-
-	hs = append(hs, cb)
+func (c *client) OnNotify(cb OnNotifyFunc) {
+	c.onNotifyHandler = cb
 }
 
 func (c *client) OnClose(cb OnCloseFunc) {
@@ -214,21 +206,24 @@ func (c *client) input() {
 	var res protocol.Response
 	var noti protocol.Notification
 	var messageType int
-	var reader io.Reader
+	var buff []byte
 
 	for err == nil {
 		res = protocol.Response{}
-		if messageType, reader, err = c.conn.NextReader(); nil != err {
-			break
+
+		if messageType, buff, err = c.conn.ReadMessage(); nil != err {
+			logging.Logger.Error(fmt.Sprintf("Client: Reader error[%v]", err))
+			continue
 		}
 		logging.Logger.Debug(fmt.Sprintf("Client: messageType:%d", messageType))
 
-		if err = json.NewDecoder(reader).Decode(res); nil != err {
+		if err = json.Unmarshal(buff, &res); nil != err {
 			noti = protocol.Notification{}
-			if err = json.NewDecoder(reader).Decode(noti); nil != err {
-				break
+			if err = json.Unmarshal(buff, &noti); nil != err {
+				logging.Logger.Error(fmt.Sprintf("Client: Decode error[%v]", err))
+				continue
 			} else {
-				err = c.onNotification(noti)
+				c.onNotification(noti)
 			}
 		} else {
 			err = c.onResponse(res)
@@ -290,17 +285,12 @@ func (c *client) onResponse(res protocol.Response) error {
 	return err
 }
 
-func (c *client) onNotification(noti protocol.Notification) error {
-	var err error
-	var hs []OnNotifyFunc
-	var ok bool
-	if hs, ok = c.onNotifyHandlers[noti.Method]; ok {
-		for _, h := range hs {
-			h(noti.Method, noti.Params)
-		}
+func (c *client) onNotification(noti protocol.Notification) {
+	if nil == c.onNotifyHandler {
+		return
 	}
 
-	return err
+	c.onNotifyHandler(noti.Method, noti.Params)
 }
 
 func (c *client) connCloseHandler(code int, text string) error {