rpc: add subscriber queue overflow check

Server-side test is added, but disabled because of its unreliability.

pkg/rpc/client/wsclient.go

@@ -125,7 +125,7 @@ readloop:
 			}
 			var slice []json.RawMessage
 			err = json.Unmarshal(rr.RawParams, &slice)
-			if err != nil || len(slice) != 1 {
+			if err != nil || (event != response.MissedEventID && len(slice) != 1) {
 				// Bad event received.
 				break
 			}
@@ -139,14 +139,18 @@ readloop:
 				val = new(result.NotificationEvent)
 			case response.ExecutionEventID:
 				val = new(result.ApplicationLog)
+			case response.MissedEventID:
+				// No value.
 			default:
 				// Bad event received.
 				break readloop
 			}
-			err = json.Unmarshal(slice[0], val)
+			if event != response.MissedEventID {
-			if err != nil || len(slice) != 1 {
+				err = json.Unmarshal(slice[0], val)
-				// Bad event received.
+				if err != nil {
-				break
+					// Bad event received.
+					break
+				}
 			}
 			c.Notifications <- Notification{event, val}
 		} else if rr.RawID != nil && (rr.Error != nil || rr.Result != nil) {

pkg/rpc/client/wsclient_test.go

@@ -108,6 +108,7 @@ func TestWSClientEvents(t *testing.T) {
 		`{"jsonrpc":"2.0","method":"notification_from_execution","params":[{"contract":"0x1b4357bff5a01bdf2a6581247cf9ed1e24629176","state":{"type":"Array","value":[{"type":"ByteArray","value":"636f6e74726163742063616c6c"},{"type":"ByteArray","value":"7472616e73666572"},{"type":"Array","value":[{"type":"ByteArray","value":"769162241eedf97c2481652adf1ba0f5bf57431b"},{"type":"ByteArray","value":"316e851039019d39dfc2c37d6c3fee19fd580987"},{"type":"Integer","value":"1000"}]}]}}]}`,
 		`{"jsonrpc":"2.0","method":"transaction_added","params":[{"txid":"0xe1cd5e57e721d2a2e05fb1f08721b12057b25ab1dd7fd0f33ee1639932fdfad7","size":277,"type":"InvocationTransaction","version":1,"nonce":9,"sender":"ALHF9wsXZVEuCGgmDA6ZNsCLtrb4A1g4yG","sys_fee":"0","net_fee":"0.0037721","valid_until_block":1200,"attributes":[],"cosigners":[{"account":"0x870958fd19ee3f6c7dc3c2df399d013910856e31","scopes":1}],"vin":[],"vout":[],"scripts":[{"invocation":"0c4027727296b84853c5d9e07fb8a40e885246ae25641383b16eefbe92027ecb1635b794aacf6bbfc3e828c73829b14791c483d19eb758b57638e3191393dbf2d288","verification":"0c2102b3622bf4017bdfe317c58aed5f4c753f206b7db896046fa7d774bbc4bf7f8dc20b410a906ad4"}],"script":"01e8030c14316e851039019d39dfc2c37d6c3fee19fd5809870c14769162241eedf97c2481652adf1ba0f5bf57431b13c00c087472616e736665720c14769162241eedf97c2481652adf1ba0f5bf57431b41627d5b5238"}]}`,
 		`{"jsonrpc":"2.0","method":"block_added","params":[{"version":0,"previousblockhash":"0x04f7580b111ec75f0ce68d3a9fd70a0544b4521b4a98541694d8575c548b759e","merkleroot":"0xb2c7230ebee4cb83bc03afadbba413e6bca8fcdeaf9c077bea060918da0e52a1","time":1590006200,"height":207,"next_consensus":"0x4138145d67638db97b402ee0b6751ef16253ecab","script":{"invocation":"0c4063429fca5ff75c964d9e38179c75978e33f8174d91a780c2e825265cf2447281594afdd5f3e216dcaf5ff0693aec83f415996cf224454495495f6bd0a4c5d08f0c4099680903a954278580d8533121c2cd3e53a089817b6a784901ec06178a60b5f1da6e70422bdcadc89029767e08d66ce4180b99334cb2d42f42e4216394af15920c4067d5e362189e48839a24e187c59d46f5d9db862c8a029777f1548b19632bfdc73ad373827ed02369f925e89c2303b64e6b9838dca229949b9b9d3bd4c0c3ed8f0c4021d4c00d4522805883f1db929554441bcbbee127c48f6b7feeeb69a72a78c7f0a75011663e239c0820ef903f36168f42936de10f0ef20681cb735a4b53d0390f","verification":"130c2102103a7f7dd016558597f7960d27c516a4394fd968b9e65155eb4b013e4040406e0c2102a7bc55fe8684e0119768d104ba30795bdcc86619e864add26156723ed185cd620c2102b3622bf4017bdfe317c58aed5f4c753f206b7db896046fa7d774bbc4bf7f8dc20c2103d90c07df63e690ce77912e10ab51acc944b66860237b608c4f8f8309e71ee699140b413073b3bb"},"consensus_data":{"primary":0,"nonce":1111},"tx":[{"txid":"0xf736cd91ab84062a21a09b424346b241987f6245ffe8c2b2db39d595c3c222f7","size":204,"type":"InvocationTransaction","version":1,"nonce":8,"sender":"ALHF9wsXZVEuCGgmDA6ZNsCLtrb4A1g4yG","sys_fee":"0","net_fee":"0.0030421","valid_until_block":1200,"attributes":[],"cosigners":[],"vin":[],"vout":[],"scripts":[{"invocation":"0c4016e7a112742409cdfaad89dcdbcb52c94c5c1a69dfe5d8b999649eaaa787e31ca496d1734d6ea606c749ad36e9a88892240ae59e0efa7f544e0692124898d512","verification":"0c2102b3622bf4017bdfe317c58aed5f4c753f206b7db896046fa7d774bbc4bf7f8dc20b410a906ad4"}],"script":"10c00c04696e69740c14769162241eedf97c2481652adf1ba0f5bf57431b41627d5b52"},{"txid":"0xe1cd5e57e721d2a2e05fb1f08721b12057b25ab1dd7fd0f33ee1639932fdfad7","size":277,"type":"InvocationTransaction","version":1,"nonce":9,"sender":"ALHF9wsXZVEuCGgmDA6ZNsCLtrb4A1g4yG","sys_fee":"0","net_fee":"0.0037721","valid_until_block":1200,"attributes":[],"cosigners":[{"account":"0x870958fd19ee3f6c7dc3c2df399d013910856e31","scopes":1}],"vin":[],"vout":[],"scripts":[{"invocation":"0c4027727296b84853c5d9e07fb8a40e885246ae25641383b16eefbe92027ecb1635b794aacf6bbfc3e828c73829b14791c483d19eb758b57638e3191393dbf2d288","verification":"0c2102b3622bf4017bdfe317c58aed5f4c753f206b7db896046fa7d774bbc4bf7f8dc20b410a906ad4"}],"script":"01e8030c14316e851039019d39dfc2c37d6c3fee19fd5809870c14769162241eedf97c2481652adf1ba0f5bf57431b13c00c087472616e736665720c14769162241eedf97c2481652adf1ba0f5bf57431b41627d5b5238"}]}]}`,
+		`{"jsonrpc":"2.0","method":"event_missed","params":[]}`,
 	}
 	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
 		if req.URL.Path == "/ws" && req.Method == "GET" {

pkg/rpc/response/events.go

@@ -23,6 +23,8 @@ const (
 	NotificationEventID
 	// ExecutionEventID is used for `transaction_executed` events.
 	ExecutionEventID
+	// MissedEventID notifies user of missed events.
+	MissedEventID EventID = 255
 )
 
 // String is a good old Stringer implementation.
@@ -36,6 +38,8 @@ func (e EventID) String() string {
 		return "notification_from_execution"
 	case ExecutionEventID:
 		return "transaction_executed"
+	case MissedEventID:
+		return "event_missed"
 	default:
 		return "unknown"
 	}
@@ -52,6 +56,8 @@ func GetEventIDFromString(s string) (EventID, error) {
 		return NotificationEventID, nil
 	case "transaction_executed":
 		return ExecutionEventID, nil
+	case "event_missed":
+		return MissedEventID, nil
 	default:
 		return 255, errors.New("invalid stream name")
 	}

pkg/rpc/server/server.go

@@ -296,37 +296,49 @@ func (s *Server) handleRequest(req *request.In, sub *subscriber) response.Raw {
 
 func (s *Server) handleWsWrites(ws *websocket.Conn, resChan <-chan response.Raw, subChan <-chan *websocket.PreparedMessage) {
 	pingTicker := time.NewTicker(wsPingPeriod)
-	defer ws.Close()
+eventloop:
-	defer pingTicker.Stop()
 	for {
 		select {
 		case <-s.shutdown:
-			// Signal to the reader routine.
+			break eventloop
-			ws.Close()
-			return
 		case event, ok := <-subChan:
 			if !ok {
-				return
+				break eventloop
 			}
 			ws.SetWriteDeadline(time.Now().Add(wsWriteLimit))
 			if err := ws.WritePreparedMessage(event); err != nil {
-				return
+				break eventloop
 			}
 		case res, ok := <-resChan:
 			if !ok {
-				return
+				break eventloop
 			}
 			ws.SetWriteDeadline(time.Now().Add(wsWriteLimit))
 			if err := ws.WriteJSON(res); err != nil {
-				return
+				break eventloop
 			}
 		case <-pingTicker.C:
 			ws.SetWriteDeadline(time.Now().Add(wsWriteLimit))
 			if err := ws.WriteMessage(websocket.PingMessage, []byte{}); err != nil {
-				return
+				break eventloop
 			}
 		}
 	}
+	ws.Close()
+	pingTicker.Stop()
+	// Drain notification channel as there might be some goroutines blocked
+	// on it.
+drainloop:
+	for {
+		select {
+		case _, ok := <-subChan:
+			if !ok {
+				break drainloop
+			}
+		default:
+			break drainloop
+		}
+	}
 }
 
 func (s *Server) handleWsReads(ws *websocket.Conn, resChan chan<- response.Raw, subscr *subscriber) {
@@ -1130,7 +1142,7 @@ func (s *Server) subscribe(reqParams request.Params, sub *subscriber) (interface
 		return nil, response.ErrInvalidParams
 	}
 	event, err := response.GetEventIDFromString(streamName)
-	if err != nil {
+	if err != nil || event == response.MissedEventID {
 		return nil, response.ErrInvalidParams
 	}
 	s.subsLock.Lock()
@@ -1232,6 +1244,20 @@ func (s *Server) unsubscribeFromChannel(event response.EventID) {
 }
 
 func (s *Server) handleSubEvents() {
+	b, err := json.Marshal(response.Notification{
+		JSONRPC: request.JSONRPCVersion,
+		Event:   response.MissedEventID,
+		Payload: make([]interface{}, 0),
+	})
+	if err != nil {
+		s.log.Error("fatal: failed to marshal overflow event", zap.Error(err))
+		return
+	}
+	overflowMsg, err := websocket.NewPreparedMessage(websocket.TextMessage, b)
+	if err != nil {
+		s.log.Error("fatal: failed to prepare overflow message", zap.Error(err))
+		return
+	}
 chloop:
 	for {
 		var resp = response.Notification{
@@ -1258,10 +1284,13 @@ chloop:
 		s.subsLock.RLock()
 	subloop:
 		for sub := range s.subscribers {
+			if sub.overflown.Load() {
+				continue
+			}
 			for _, subID := range sub.feeds {
 				if subID == resp.Event {
 					if msg == nil {
-						b, err := json.Marshal(resp)
+						b, err = json.Marshal(resp)
 						if err != nil {
 							s.log.Error("failed to marshal notification",
 								zap.Error(err),
@@ -1276,7 +1305,16 @@ chloop:
 							break subloop
 						}
 					}
-					sub.writer <- msg
+					select {
+					case sub.writer <- msg:
+					default:
+						sub.overflown.Store(true)
+						// MissedEvent is to be delivered eventually.
+						go func(sub *subscriber) {
+							sub.writer <- overflowMsg
+							sub.overflown.Store(false)
+						}(sub)
+					}
 					// The message is sent only once per subscriber.
 					break
 				}

pkg/rpc/server/subscription.go

@@ -3,14 +3,15 @@ package server
 import (
 	"github.com/gorilla/websocket"
 	"github.com/nspcc-dev/neo-go/pkg/rpc/response"
+	"go.uber.org/atomic"
 )
 
 type (
 	// subscriber is an event subscriber.
 	subscriber struct {
-		writer chan<- *websocket.PreparedMessage
+		writer    chan<- *websocket.PreparedMessage
-		ws     *websocket.Conn
+		ws        *websocket.Conn
-
+		overflown atomic.Bool
 		// These work like slots as there is not a lot of them (it's
 		// cheaper doing it this way rather than creating a map),
 		// pointing to EventID is an obvious overkill at the moment, but

pkg/rpc/server/subscription_test.go

@@ -160,6 +160,7 @@ func TestBadSubUnsub(t *testing.T) {
 		"no params":              `{"jsonrpc": "2.0", "method": "subscribe", "params": [], "id": 1}`,
 		"bad (non-string) event": `{"jsonrpc": "2.0", "method": "subscribe", "params": [1], "id": 1}`,
 		"bad (wrong) event":      `{"jsonrpc": "2.0", "method": "subscribe", "params": ["block_removed"], "id": 1}`,
+		"missed event":           `{"jsonrpc": "2.0", "method": "subscribe", "params": ["event_missed"], "id": 1}`,
 	}
 	var unsubCases = map[string]string{
 		"no params":         `{"jsonrpc": "2.0", "method": "unsubscribe", "params": [], "id": 1}`,
@@ -225,3 +226,42 @@ func TestWSClientsLimit(t *testing.T) {
 		doSomeWSRequest(t, wss[i])
 	}
 }
+
+// The purpose of this test is to overflow buffers on server side to
+// receive a 'missed' event. But it's actually hard to tell when exactly
+// that's going to happen because of network-level buffering, typical
+// number seen in tests is around ~3500 events, but it's not reliable enough,
+// thus this test is disabled.
+func testSubscriptionOverflow(t *testing.T) {
+	const blockCnt = notificationBufSize * 5
+	var receivedMiss bool
+
+	chain, rpcSrv, c, respMsgs, finishedFlag := initCleanServerAndWSClient(t)
+
+	defer chain.Close()
+	defer rpcSrv.Shutdown()
+
+	resp := callWSGetRaw(t, c, `{"jsonrpc": "2.0","method": "subscribe","params": ["block_added"],"id": 1}`, respMsgs)
+	require.Nil(t, resp.Error)
+	require.NotNil(t, resp.Result)
+
+	// Push a lot of new blocks, but don't read events for them.
+	for i := 0; i < blockCnt; i++ {
+		b := newBlock(t, chain, 1)
+		require.NoError(t, chain.AddBlock(b))
+	}
+	for i := 0; i < blockCnt; i++ {
+		resp := getNotification(t, respMsgs)
+		if resp.Event != response.BlockEventID {
+			require.Equal(t, response.MissedEventID, resp.Event)
+			receivedMiss = true
+			break
+		}
+	}
+	require.Equal(t, true, receivedMiss)
+	// `Missed` is the last event and there is nothing afterwards.
+	require.Equal(t, 0, len(respMsgs))
+
+	finishedFlag.CAS(false, true)
+	c.Close()
+}