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,18 +54,25 @@ 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
+	blockSubChannels   map[int]byte
+	storageSubChannels map[int]byte
+	qtyListeners       int
+	subscriptions      map[int]Listener
+	storageAPI         modules.StorageAPI
+	blockAPI           modules.BlockAPI
 }
 
+var logger log.Logger
+
 // NewHTTPServer creates a new http server and registers an associated rpc server
 func NewHTTPServer(cfg *HTTPServerConfig) *HTTPServer {
-	logger := log.New("pkg", "rpc")
+	logger = log.New("pkg", "rpc")
 	h := log.StreamHandler(os.Stdout, log.TerminalFormat())
 	logger.SetHandler(log.LvlFilterHandler(cfg.LogLvl, h))
 
@@ -77,10 +82,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 +152,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{} {