This commit is contained in:
Anton Kaliaev
2019-01-30 20:32:53 +04:00
parent 641182e5d3
commit 54cc5100f8
21 changed files with 170 additions and 212 deletions

View File

@ -2,7 +2,6 @@ package rpcserver
import (
"bytes"
"context"
"encoding/hex"
"encoding/json"
"fmt"
@ -434,8 +433,8 @@ type wsConnection struct {
// Send pings to server with this period. Must be less than readWait, but greater than zero.
pingPeriod time.Duration
// object that is used to subscribe / unsubscribe from events
eventSub types.EventSubscriber
// see DisconnectCallback option.
disconnectCallback func(remoteAddr string)
}
// NewWSConnection wraps websocket.Conn.
@ -468,12 +467,11 @@ func NewWSConnection(
return wsc
}
// EventSubscriber sets object that is used to subscribe / unsubscribe from
// events - not Goroutine-safe. If none given, default node's eventBus will be
// used.
func EventSubscriber(eventSub types.EventSubscriber) func(*wsConnection) {
// DisconnectCallback can be used optionally to set a callback, which will be
// called upon disconnect - not Goroutine-safe.
func DisconnectCallback(cb func(remoteAddr string)) func(*wsConnection) {
return func(wsc *wsConnection) {
wsc.eventSub = eventSub
wsc.disconnectCallback = cb
}
}
@ -526,8 +524,8 @@ func (wsc *wsConnection) OnStart() error {
func (wsc *wsConnection) OnStop() {
// Both read and write loops close the websocket connection when they exit their loops.
// The writeChan is never closed, to allow WriteRPCResponse() to fail.
if wsc.eventSub != nil {
wsc.eventSub.UnsubscribeAll(context.TODO(), wsc.remoteAddr)
if wsc.disconnectCallback != nil {
wsc.disconnectCallback(wsc.remoteAddr)
}
}
@ -537,11 +535,6 @@ func (wsc *wsConnection) GetRemoteAddr() string {
return wsc.remoteAddr
}
// GetEventSubscriber implements WSRPCConnection by returning event subscriber.
func (wsc *wsConnection) GetEventSubscriber() types.EventSubscriber {
return wsc.eventSub
}
// WriteRPCResponse pushes a response to the writeChan, and blocks until it is accepted.
// It implements WSRPCConnection. It is Goroutine-safe.
func (wsc *wsConnection) WriteRPCResponse(resp types.RPCResponse) {