dot/rpc implement state_subscribeStorage RPC WebSocket method

dot/rpc/http.go

@@ -20,23 +20,21 @@ import (
 	"fmt"
 	"net/http"
 	"os"
+	"sync"
 
 	"github.com/ChainSafe/gossamer/dot/rpc/modules"
-	"github.com/ChainSafe/gossamer/dot/types"
+	log "github.com/ChainSafe/log15"
 	"github.com/gorilla/mux"
 	"github.com/gorilla/rpc/v2"
 	"github.com/gorilla/websocket"
-
-	log "github.com/ChainSafe/log15"
 )
 
 // HTTPServer gateway for RPC server
 type HTTPServer struct {
 	logger       log.Logger
 	rpcServer    *rpc.Server // Actual RPC call handler
 	serverConfig *HTTPServerConfig
-	blockChan    chan *types.Block
-	chanID       byte // channel ID
+	wsConns      []*WSConn
 }
 
 // HTTPServerConfig configures the HTTPServer
@@ -56,13 +54,18 @@ type HTTPServerConfig struct {
 	WSEnabled           bool
 	WSPort              uint32
 	Modules             []string
-	WSSubscriptions     map[uint32]*WebSocketSubscription
 }
 
-// WebSocketSubscription holds subscription details
-type WebSocketSubscription struct {
-	WSConnection     *websocket.Conn
-	SubscriptionType int
+// WSConn struct to hold WebSocket Connection references
+type WSConn struct {
+	wsconn             *websocket.Conn
+	mu                 sync.Mutex
+	serverConfig       *HTTPServerConfig
+	logger             log.Logger
+	blockSubChannels   map[int]byte
+	storageSubChannels map[int]byte
+	qtyListeners       int
+	subscriptions      map[int]Listener
 }
 
 // NewHTTPServer creates a new http server and registers an associated rpc server
@@ -77,10 +80,6 @@ func NewHTTPServer(cfg *HTTPServerConfig) *HTTPServer {
 		serverConfig: cfg,
 	}
 
-	if cfg.WSSubscriptions == nil {
-		cfg.WSSubscriptions = make(map[uint32]*WebSocketSubscription)
-	}
-
 	server.RegisterModules(cfg.Modules)
 	return server
 }
@@ -151,25 +150,30 @@ func (h *HTTPServer) Start() error {
 		}
 	}()
 
-	// init and start block received listener routine
-	if h.serverConfig.BlockAPI != nil {
-		var err error
-		h.blockChan = make(chan *types.Block)
-		h.chanID, err = h.serverConfig.BlockAPI.RegisterImportedChannel(h.blockChan)
-		if err != nil {
-			return err
-		}
-		go h.blockReceivedListener()
-	}
-
 	return nil
 }
 
 // Stop stops the server
 func (h *HTTPServer) Stop() error {
 	if h.serverConfig.WSEnabled {
-		h.serverConfig.BlockAPI.UnregisterImportedChannel(h.chanID)
-		close(h.blockChan)
+		// close all channels and websocket connections
+		for _, conn := range h.wsConns {
+			for _, sub := range conn.subscriptions {
+				switch v := sub.(type) {
+				case *StorageChangeListener:
+					h.serverConfig.StorageAPI.UnregisterStorageChangeChannel(v.chanID)
+					close(v.channel)
+				case *BlockListener:
+					h.serverConfig.BlockAPI.UnregisterImportedChannel(v.chanID)
+					close(v.channel)
+				}
+			}
+
+			err := conn.wsconn.Close()
+			if err != nil {
+				h.logger.Error("error closing websocket connection", "error", err)
+			}
+		}
 	}
 	return nil
 }

dot/rpc/modules/api.go

@@ -3,6 +3,7 @@ package modules
 import (
 	"math/big"
 
+	"github.com/ChainSafe/gossamer/dot/state"
 	"github.com/ChainSafe/gossamer/dot/types"
 	"github.com/ChainSafe/gossamer/lib/common"
 	"github.com/ChainSafe/gossamer/lib/crypto"
@@ -14,6 +15,8 @@ import (
 type StorageAPI interface {
 	GetStorage(key []byte) ([]byte, error)
 	Entries() map[string][]byte
+	RegisterStorageChangeChannel(ch chan<- *state.KeyValue) (byte, error)
+	UnregisterStorageChangeChannel(id byte)
 }
 
 // BlockAPI is the interface for the block state

dot/rpc/modules/state.go

@@ -285,9 +285,11 @@ func (sm *StateModule) SubscribeRuntimeVersion(r *http.Request, req *StateStorag
 	return sm.GetRuntimeVersion(r, nil, res)
 }
 
-// SubscribeStorage isn't implemented properly yet.
-func (sm *StateModule) SubscribeStorage(r *http.Request, req *StateStorageQueryRangeRequest, res *StorageChangeSetResponse) {
-	// TODO implement change storage trie so that block hash parameter works (See issue #834)
+// SubscribeStorage Storage subscription. If storage keys are specified, it creates a message for each block which
+//  changes the specified storage keys. If none are specified, then it creates a message for every block.
+//  This endpoint communicates over the Websocket protocol, but this func should remain here so it's added to rpc_methods list
+func (sm *StateModule) SubscribeStorage(r *http.Request, req *StateStorageQueryRangeRequest, res *StorageChangeSetResponse) error {
+	return nil
 }
 
 func convertAPIs(in []*runtime.API_Item) []interface{} {