network: Kademlia Load Balancing

network/discovery.go

@@ -21,6 +21,7 @@ import (
 	"fmt"
 	"sync"
 
+	"github.com/ethereum/go-ethereum/common/hexutil"
 	"github.com/ethereum/go-ethereum/rlp"
 	"github.com/ethersphere/swarm/pot"
 )
@@ -37,6 +38,7 @@ type Peer struct {
 	mtx       sync.RWMutex    //
 	peers     map[string]bool // tracks node records sent to the peer
 	depth     uint8           // the proximity order advertised by remote as depth of saturation
+	key       string          // peer key. Hex form of Address()
 }
 
 // NewPeer constructs a discovery peer
@@ -45,6 +47,7 @@ func NewPeer(p *BzzPeer, kad *Kademlia) *Peer {
 		kad:     kad,
 		BzzPeer: p,
 		peers:   make(map[string]bool),
+		key:     hexutil.Encode(p.Address()),
 	}
 	// record remote as seen so we never send a peer its own record
 	d.seen(p.BzzAddr)
@@ -66,6 +69,16 @@ func (d *Peer) HandleMsg(ctx context.Context, msg interface{}) error {
 	}
 }
 
+// Key returns a string representation of this peer to be used in maps.
+func (d *Peer) Key() string {
+	return d.key
+}
+
+// Label returns a short string representation for debugging purposes
+func (d *Peer) Label() string {
+	return d.key[:4]
+}
+
 // NotifyDepth sends a message to all connections if depth of saturation is changed
 func NotifyDepth(depth uint8, kad *Kademlia) {
 	f := func(val *Peer, po int) bool {

network/gopubsub/pubsub.go

@@ -0,0 +1,126 @@
+package gopubsub
+
+import (
+	"fmt"
+	"strconv"
+	"sync"
+)
+
+//PubSubChannel represents a pubsub system where subscriber can .Subscribe() and publishers can .Publish() or .Close().
+type PubSubChannel struct {
+	subscriptions []*Subscription
+	subsMutex     sync.RWMutex
+	nextId        int
+}
+
+//Subscription is created in PubSubChannel using pubSub.Subscribe(). Subscribers can receive using .ReceiveChannel().
+// or .Unsubscribe()
+type Subscription struct {
+	closed    bool
+	removeSub func()
+	signal    chan interface{}
+	closeOnce sync.Once
+	id        string
+}
+
+//New creates a new PubSubChannel.
+func New() *PubSubChannel {
+	return &PubSubChannel{
+		subscriptions: make([]*Subscription, 0),
+	}
+}
+
+//Subscribe creates a subscription to a channel, each subscriber should keep its own Subscription instance.
+func (psc *PubSubChannel) Subscribe() *Subscription {
+	psc.subsMutex.Lock()
+	defer psc.subsMutex.Unlock()
+	newSubscription := newSubscription(strconv.Itoa(psc.nextId))
+	psc.nextId++
+	psc.subscriptions = append(psc.subscriptions, &newSubscription)
+	newSubscription.removeSub = func() {
+		psc.subsMutex.Lock()
+		defer psc.subsMutex.Unlock()
+
+		for i, subscription := range psc.subscriptions {
+			if subscription.signal == newSubscription.signal {
+				fmt.Println("Unsubscribing", "id", subscription.id)
+				subscription.closed = true
+				psc.subscriptions = append(psc.subscriptions[:i], psc.subscriptions[i+1:]...)
+			}
+		}
+	}
+	return &newSubscription
+}
+
+//Publish broadcasts a message asynchronously to each subscriber.
+//If some of the subscriptions(channels) has been marked as closeable, it does it now.
+func (psc *PubSubChannel) Publish(msg interface{}) {
+	psc.subsMutex.RLock()
+	defer psc.subsMutex.RUnlock()
+	for i, sub := range psc.subscriptions {
+		if sub.closed {
+			fmt.Println("Subscription was closed", "id", sub.id)
+			sub.closeChannel()
+		} else {
+			go func(sub *Subscription, index int) {
+				sub.signal <- msg
+			}(sub, i)
+
+		}
+	}
+}
+
+//NumSubscriptions returns how many subscriptions are currently active.
+func (psc *PubSubChannel) NumSubscriptions() int {
+	psc.subsMutex.RLock()
+	defer psc.subsMutex.RUnlock()
+	return len(psc.subscriptions)
+}
+
+//Close cancels all subscriptions closing the channels associated with them.
+//Usually the publisher is in charge of calling Close().
+func (psc *PubSubChannel) Close() {
+	psc.subsMutex.Lock()
+	defer psc.subsMutex.Unlock()
+	for _, sub := range psc.subscriptions {
+		sub.closed = true
+		sub.closeChannel()
+	}
+}
+
+//Unsubscribe cancels subscription from the subscriber side. Channel is marked as closed but only writer should close it.
+func (sub *Subscription) Unsubscribe() {
+	sub.closed = true
+	sub.removeSub()
+}
+
+//ReceiveChannel returns the channel where the subscriber will receive messages.
+func (sub *Subscription) ReceiveChannel() <-chan interface{} {
+	return sub.signal
+}
+
+//IsClosed returns if the subscription is closed via Unsubscribe() or Close() in the pubSub that creates it.
+func (sub *Subscription) IsClosed() bool {
+	return sub.closed
+}
+
+//ID returns a unique id in the PubSubChannel of this subscription. Useful for debugging.
+func (sub *Subscription) ID() string {
+	return sub.id
+}
+
+func (sub *Subscription) closeChannel() {
+	sub.closeOnce.Do(func() {
+		close(sub.signal)
+	})
+}
+
+func newSubscription(id string) Subscription {
+	return Subscription{
+		closed:    false,
+		removeSub: nil,
+		signal:    make(chan interface{}),
+		closeOnce: sync.Once{},
+		id:        id,
+	}
+}