[dot/state] implement state_subscribeStorage

dot/rpc/modules/api.go

@@ -16,7 +16,7 @@ type StorageAPI interface {
 	GetStorage(root *common.Hash, key []byte) ([]byte, error)
 	GetStorageByBlockHash(bhash common.Hash, key []byte) ([]byte, error)
 	Entries(root *common.Hash) (map[string][]byte, error)
-	RegisterStorageChangeChannel(ch chan<- *state.KeyValue) (byte, error)
+	RegisterStorageChangeChannel(sub state.StorageSubscription) (byte, error)
 	UnregisterStorageChangeChannel(id byte)
 	GetStateRootFromBlock(bhash *common.Hash) (*common.Hash, error)
 }

dot/rpc/websocket.go

@@ -182,7 +182,7 @@ func (c *WSConn) handleComm() {
 			case "state_subscribeStorage":
 				scl, err2 := c.initStorageChangeListener(reqid, params)
 				if err2 != nil {
-					logger.Warn("failed to create state change listener", "error", err)
+					logger.Warn("failed to create state change listener", "error", err2)
 					continue
 				}
 				c.startListener(scl)
@@ -256,38 +256,45 @@ type Listener interface {
 
 // StorageChangeListener for listening to state change channels
 type StorageChangeListener struct {
-	channel chan *state.KeyValue
-	filter  map[string]bool
+	channel chan *state.SubscriptionResult
 	wsconn  *WSConn
 	chanID  byte
 	subID   int
 }
 
 func (c *WSConn) initStorageChangeListener(reqID float64, params interface{}) (int, error) {
 	scl := &StorageChangeListener{
-		channel: make(chan *state.KeyValue),
-		filter:  make(map[string]bool),
+		channel: make(chan *state.SubscriptionResult),
 		wsconn:  c,
 	}
+	sub := &state.StorageSubscription{
+		Filter:   make(map[string]bool),
+		Listener: scl.channel,
+	}
+
 	pA := params.([]interface{})
 	for _, param := range pA {
-		switch param.(type) {
+		switch p := param.(type) {
 		case []interface{}:
-			for _, p := range param.([]interface{}) {
-				scl.filter[p.(string)] = true
+			for _, pp := range param.([]interface{}) {
+				sub.Filter[pp.(string)] = true
 			}
+		case string:
+			sub.Filter[p] = true
 		default:
 			return 0, fmt.Errorf("unknow parameter type")
 		}
 	}
+
 	if c.storageAPI == nil {
 		err := c.safeSendError(reqID, nil, "error StorageAPI not set")
 		if err != nil {
 			logger.Warn("error sending error message", "error", err)
 		}
 		return 0, fmt.Errorf("error StorageAPI not set")
 	}
-	chanID, err := c.storageAPI.RegisterStorageChangeChannel(scl.channel)
+
+	chanID, err := c.storageAPI.RegisterStorageChangeChannel(*sub)
 	if err != nil {
 		return 0, err
 	}
@@ -313,14 +320,17 @@ func (l *StorageChangeListener) Listen() {
 			continue
 		}
 
-		//check if change key is in subscription filter
-		cKey := common.BytesToHex(change.Key)
-		if len(l.filter) > 0 && !l.filter[cKey] {
-			continue
+		result := make(map[string]interface{})
+		result["block"] = change.Hash.String()
+		changes := [][]string{}
+		for _, v := range change.Changes {
+			kv := []string{common.BytesToHex(v.Key), common.BytesToHex(v.Value)}
+			changes = append(changes, kv)
 		}
+		result["changes"] = changes
 
 		changeM := make(map[string]interface{})
-		changeM["result"] = []string{cKey, common.BytesToHex(change.Value)}
+		changeM["result"] = result
 		changeM["subscription"] = l.subID
 		res := newSubcriptionBaseResponseJSON()
 		res.Method = "state_storage"
@@ -329,7 +339,6 @@ func (l *StorageChangeListener) Listen() {
 		if err != nil {
 			logger.Error("error sending websocket message", "error", err)
 		}
-
 	}
 }
 

dot/rpc/websocket_test.go

@@ -131,7 +131,7 @@ func (m *MockStorageAPI) Entries(_ *common.Hash) (map[string][]byte, error) {
 func (m *MockStorageAPI) GetStorageByBlockHash(_ common.Hash, key []byte) ([]byte, error) {
 	return nil, nil
 }
-func (m *MockStorageAPI) RegisterStorageChangeChannel(ch chan<- *state.KeyValue) (byte, error) {
+func (m *MockStorageAPI) RegisterStorageChangeChannel(sub state.StorageSubscription) (byte, error) {
 	return 0, nil
 }
 func (m *MockStorageAPI) UnregisterStorageChangeChannel(id byte) {

dot/state/storage.go

@@ -49,8 +49,8 @@ type StorageState struct {
 	lock   sync.RWMutex
 
 	// change notifiers
-	changed     map[byte]chan<- *KeyValue
-	changedLock sync.RWMutex
+	changedLock   sync.RWMutex
+	subscriptions map[byte]*StorageSubscription
 }
 
 // NewStorageState creates a new StorageState backed by the given trie and database located at basePath.
@@ -67,11 +67,11 @@ func NewStorageState(db chaindb.Database, blockState *BlockState, t *trie.Trie)
 	tries[t.MustHash()] = t
 
 	return &StorageState{
-		blockState: blockState,
-		tries:      tries,
-		baseDB:     db,
-		db:         chaindb.NewTable(db, storagePrefix),
-		changed:    make(map[byte]chan<- *KeyValue),
+		blockState:    blockState,
+		tries:         tries,
+		baseDB:        db,
+		db:            chaindb.NewTable(db, storagePrefix),
+		subscriptions: make(map[byte]*StorageSubscription),
 	}, nil
 }
 
@@ -162,7 +162,49 @@ func (s *StorageState) StoreInDB(root common.Hash) error {
 		return errTrieDoesNotExist(root)
 	}
 
-	return StoreTrie(s.baseDB, s.tries[root])
+	err := StoreTrie(s.baseDB, s.tries[root])
+	if err != nil {
+		return err
+	}
+
+	// notify subscribers of database changes
+	for _, sub := range s.subscriptions {
+		subRes := &SubscriptionResult{
+			Hash: root,
+		}
+		if len(sub.Filter) == 0 {
+			// no filter, so send all changes
+			ent := s.tries[root].Entries()
+			for k, v := range ent {
+				if k != ":code" {
+					// todo, currently we're ignoring :code since this is a lot of data
+					kv := &KeyValue{
+						Key:   common.MustHexToBytes(fmt.Sprintf("0x%x", k)),
+						Value: v,
+					}
+					subRes.Changes = append(subRes.Changes, *kv)
+				}
+			}
+
+		} else {
+			// filter result to include only interested keys
+			for k := range sub.Filter {
+				value, err := s.tries[root].Get(common.MustHexToBytes(k))
+				if err != nil {
+					logger.Error("Error retrieving value from state tries")
+					continue
+				}
+				kv := &KeyValue{
+					Key:   common.MustHexToBytes(k),
+					Value: value,
+				}
+				subRes.Changes = append(subRes.Changes, *kv)
+			}
+		}
+		s.notifyChanged(subRes)
+	}
+
+	return nil
 }
 
 // LoadFromDB loads an encoded trie from the DB where the key is `root`
@@ -353,48 +395,6 @@ func (s *StorageState) GetBalance(hash *common.Hash, key [32]byte) (uint64, erro
 	return binary.LittleEndian.Uint64(bal), nil
 }
 
-// setStorage set the storage value for a given key in the trie. only for testing
-func (s *StorageState) setStorage(hash *common.Hash, key []byte, value []byte) error {
-	if hash == nil {
-		sr, err := s.blockState.BestBlockStateRoot()
-		if err != nil {
-			return err
-		}
-		hash = &sr
-	}
-
-	s.lock.Lock()
-	defer s.lock.Unlock()
-	kv := &KeyValue{
-		Key:   key,
-		Value: value,
-	}
-
-	if s.tries[*hash] == nil {
-		return errTrieDoesNotExist(*hash)
-	}
-
-	err := s.tries[*hash].Put(key, value)
-	if err != nil {
-		return err
-	}
-	s.notifyChanged(kv) // TODO: what is this used for? needs to be updated to work with new StorageState/TrieState API
-	return nil
-}
-
-// setBalance sets the balance for an account with the given public key. only for testing
-func (s *StorageState) setBalance(hash *common.Hash, key [32]byte, balance uint64) error {
-	skey, err := common.BalanceKey(key)
-	if err != nil {
-		return err
-	}
-
-	bb := make([]byte, 8)
-	binary.LittleEndian.PutUint64(bb, balance)
-
-	return s.setStorage(hash, skey, bb)
-}
-
 func (s *StorageState) pruneStorage(closeCh chan interface{}) {
 	for {
 		select {

dot/state/storage_notify.go

@@ -17,6 +17,8 @@ package state
 
 import (
 	"errors"
+
+	"github.com/ChainSafe/gossamer/lib/common"
 )
 
 // KeyValue struct to hold key value pairs
@@ -25,26 +27,38 @@ type KeyValue struct {
 	Value []byte
 }
 
+//SubscriptionResult holds results of storage changes
+type SubscriptionResult struct {
+	Hash    common.Hash
+	Changes []KeyValue
+}
+
+//StorageSubscription holds data for Subscription to Storage
+type StorageSubscription struct {
+	Filter   map[string]bool
+	Listener chan<- *SubscriptionResult
+}
+
 // RegisterStorageChangeChannel function to register storage change channels
-func (s *StorageState) RegisterStorageChangeChannel(ch chan<- *KeyValue) (byte, error) {
+func (s *StorageState) RegisterStorageChangeChannel(sub StorageSubscription) (byte, error) {
 	s.changedLock.RLock()
 
-	if len(s.changed) == 256 {
-		return 0, errors.New("channel limit reached")
+	if len(s.subscriptions) == 256 {
+		return 0, errors.New("storage subscriptions limit reached")
 	}
 
 	var id byte
 	for {
 		id = generateID()
-		if s.changed[id] == nil {
+		if s.subscriptions[id] == nil {
 			break
 		}
 	}
 
 	s.changedLock.RUnlock()
 
 	s.changedLock.Lock()
-	s.changed[id] = ch
+	s.subscriptions[id] = &sub
 	s.changedLock.Unlock()
 	return id, nil
 }
@@ -55,22 +69,22 @@ func (s *StorageState) UnregisterStorageChangeChannel(id byte) {
 	s.changedLock.Lock()
 	defer s.changedLock.Unlock()
 
-	delete(s.changed, id)
+	delete(s.subscriptions, id)
 }
 
-func (s *StorageState) notifyChanged(change *KeyValue) {
+func (s *StorageState) notifyChanged(change *SubscriptionResult) {
 	s.changedLock.RLock()
 	defer s.changedLock.RUnlock()
 
-	if len(s.changed) == 0 {
+	if len(s.subscriptions) == 0 {
 		return
 	}
 
-	logger.Trace("notifying changed storage chans...", "chans", s.changed)
+	logger.Trace("notifying changed storage chans...", "chans", s.subscriptions)
 
-	for _, ch := range s.changed {
-		go func(ch chan<- *KeyValue) {
+	for _, ch := range s.subscriptions {
+		go func(ch chan<- *SubscriptionResult) {
 			ch <- change
-		}(ch)
+		}(ch.Listener)
 	}
 }