swarm/network/stream: added RPC call to request subscription info

swarm/network/stream/common_test.go

@@ -46,6 +46,7 @@ var (
 	loglevel     = flag.Int("loglevel", 2, "verbosity of logs")
 	nodes        = flag.Int("nodes", 0, "number of nodes")
 	chunks       = flag.Int("chunks", 0, "number of chunks")
+	printstats   = flag.Bool("printstats", false, "print results to STDOUT")
 	useMockStore = flag.Bool("mockstore", false, "disabled mock store (default: enabled)")
 	longrunning  = flag.Bool("longrunning", false, "do run long-running tests")
 

swarm/network/stream/stream.go

@@ -838,3 +838,27 @@ func (api *API) SubscribeStream(peerId enode.ID, s Stream, history *Range, prior
 func (api *API) UnsubscribeStream(peerId enode.ID, s Stream) error {
 	return api.streamer.Unsubscribe(peerId, s)
 }
+
+/*
+GetPeerSubscriptions is a API function which allows to query a peer for stream subscriptions it has.
+It can be called via RPC.
+It returns a map of node IDs with an array of string representations of Stream objects.
+*/
+func (api *API) GetPeerSubscriptions() map[string][]string {
+	//create the empty map
+	pstreams := make(map[string][]string)
+	//iterate all streamer peers
+	for id, p := range api.streamer.peers {
+		streams := make([]string, 0)
+		//every peer has a map of stream servers
+		//every stream server represents a subscription
+		for s := range p.servers {
+			//append the string representation of the stream
+			//to the list for this peer
+			streams = append(streams, s.String())
+		}
+		//set the array of stream servers to the map
+		pstreams[id.String()] = streams
+	}
+	return pstreams
+}

swarm/network/stream/streamer_test.go

@@ -20,12 +20,21 @@ import (
 	"bytes"
 	"context"
 	"errors"
+	"fmt"
+	"os"
 	"strconv"
+	"sync"
 	"testing"
 	"time"
 
 	"github.com/ethereum/go-ethereum/crypto/sha3"
+	"github.com/ethereum/go-ethereum/node"
+	"github.com/ethereum/go-ethereum/p2p/simulations/adapters"
 	p2ptest "github.com/ethereum/go-ethereum/p2p/testing"
+	"github.com/ethereum/go-ethereum/swarm/network"
+	"github.com/ethereum/go-ethereum/swarm/network/simulation"
+	"github.com/ethereum/go-ethereum/swarm/state"
+	"github.com/ethereum/go-ethereum/swarm/storage"
 )
 
 func TestStreamerSubscribe(t *testing.T) {
@@ -921,3 +930,112 @@ func TestMaxPeerServersWithoutUnsubscribe(t *testing.T) {
 		}
 	}
 }
+
+/*
+TestGetSubscriptionsRPC sets up a simulation network of 16 nodes,
+starts the simulation, waits for SyncUpdateDelay in order to kick off
+stream registration, then tests that there are subscriptions.
+If provided with the `-printstats = true` option, it will print
+the information of who is subscribed to who to STDOUT
+*/
+func TestGetSubscriptionsRPC(t *testing.T) {
+	//arbitrarily set to 16
+	nodeCount := 16
+	//set the syncUpdateDelay for sync registrations to start
+	syncUpdateDelay := 500 * time.Millisecond
+	//create a standard sim
+	sim := simulation.New(map[string]simulation.ServiceFunc{
+		"streamer": func(ctx *adapters.ServiceContext, bucket *sync.Map) (s node.Service, cleanup func(), err error) {
+			n := ctx.Config.Node()
+			addr := network.NewAddr(n)
+			store, datadir, err := createTestLocalStorageForID(n.ID(), addr)
+			if err != nil {
+				return nil, nil, err
+			}
+			localStore := store.(*storage.LocalStore)
+			netStore, err := storage.NewNetStore(localStore, nil)
+			if err != nil {
+				return nil, nil, err
+			}
+			kad := network.NewKademlia(addr.Over(), network.NewKadParams())
+			delivery := NewDelivery(kad, netStore)
+			netStore.NewNetFetcherFunc = network.NewFetcherFactory(dummyRequestFromPeers, true).New
+			//configure so that sync registrations actually happen
+			r := NewRegistry(addr.ID(), delivery, netStore, state.NewInmemoryStore(), &RegistryOptions{
+				Retrieval:       RetrievalEnabled,
+				Syncing:         SyncingAutoSubscribe, //enable sync registrations
+				SyncUpdateDelay: syncUpdateDelay,
+			}, nil)
+
+			bucket.Store(bucketKeyRegistry, r)
+			cleanup = func() {
+				os.RemoveAll(datadir)
+				netStore.Close()
+				r.Close()
+			}
+
+			return r, cleanup, nil
+
+		},
+	})
+	defer sim.Close()
+
+	ctx, cancelSimRun := context.WithTimeout(context.Background(), 1*time.Minute)
+	defer cancelSimRun()
+
+	//upload a snapshot
+	err := sim.UploadSnapshot(fmt.Sprintf("testing/snapshot_%d.json", nodeCount))
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	//wait till healthy
+	if _, err := sim.WaitTillHealthy(ctx, 2); err != nil {
+		t.Fatal(err)
+	}
+
+	result := sim.Run(ctx, func(ctx context.Context, sim *simulation.Simulation) error {
+		//we need to wait for some time until registrations are finished...
+		time.Sleep(syncUpdateDelay + 1*time.Second)
+		nodes := sim.Net.Nodes
+
+		//iterate all nodes
+		for _, node := range nodes {
+			//create rpc client
+			client, err := node.Client()
+			if err != nil {
+				t.Fatalf("create node 1 rpc client fail: %v", err)
+			}
+
+			item, ok := sim.NodeItem(node.ID(), bucketKeyRegistry)
+			if !ok {
+				return fmt.Errorf("No registry")
+			}
+			registry := item.(*Registry)
+			//ask it for subscriptions
+			pstreams := make(map[string][]string)
+			err = client.Call(&pstreams, "stream_getPeerSubscriptions")
+			if err != nil {
+				t.Fatal(err)
+			}
+			//lenght of the subscriptions can not be smaller than number of peers
+			if len(pstreams) < len(registry.peers) {
+				t.Fatal("No subscriptions have been made")
+			}
+			//if enabled, print stats to STDOUT
+			if *printstats {
+				fmt.Println(fmt.Sprintf("node %s subscriptions:", node.String()))
+				for p, ps := range pstreams {
+					fmt.Println(fmt.Sprintf("...with node %s: ", p))
+					for _, s := range ps {
+						fmt.Println(fmt.Sprintf("......%s", s))
+					}
+				}
+			}
+		}
+		return nil
+	})
+	if result.Error != nil {
+		t.Fatal(result.Error)
+	}
+}