more fixes

ethereum/rpc/pubsub/pubsub.go

@@ -0,0 +1,128 @@
+package pubsub
+
+import (
+	"sync"
+
+	"github.com/pkg/errors"
+
+	coretypes "github.com/tendermint/tendermint/rpc/core/types"
+)
+
+type EventBus interface {
+	AddTopic(name string, src <-chan coretypes.ResultEvent) error
+	RemoveTopic(name string)
+	Subscribe(name string) (<-chan coretypes.ResultEvent, error)
+	Topics() []string
+}
+
+type memEventBus struct {
+	topics         map[string]<-chan coretypes.ResultEvent
+	topicsMux      *sync.RWMutex
+	subscribers    map[string][]chan<- coretypes.ResultEvent
+	subscribersMux *sync.RWMutex
+}
+
+func NewEventBus() EventBus {
+	return &memEventBus{
+		topics:         make(map[string]<-chan coretypes.ResultEvent),
+		topicsMux:      new(sync.RWMutex),
+		subscribers:    make(map[string][]chan<- coretypes.ResultEvent),
+		subscribersMux: new(sync.RWMutex),
+	}
+}
+
+func (m *memEventBus) Topics() (topics []string) {
+	m.topicsMux.RLock()
+	defer m.topicsMux.RUnlock()
+
+	topics = make([]string, 0, len(m.topics))
+	for topicName := range m.topics {
+		topics = append(topics, topicName)
+	}
+
+	return topics
+}
+
+func (m *memEventBus) AddTopic(name string, src <-chan coretypes.ResultEvent) error {
+	m.topicsMux.RLock()
+	_, ok := m.topics[name]
+	m.topicsMux.RUnlock()
+
+	if ok {
+		return errors.New("topic already registered")
+	}
+
+	m.topicsMux.Lock()
+	m.topics[name] = src
+	m.topicsMux.Unlock()
+
+	go m.publishTopic(name, src)
+
+	return nil
+}
+
+func (m *memEventBus) RemoveTopic(name string) {
+	m.topicsMux.Lock()
+	delete(m.topics, name)
+	m.topicsMux.Unlock()
+}
+
+func (m *memEventBus) Subscribe(name string) (<-chan coretypes.ResultEvent, error) {
+	m.topicsMux.RLock()
+	_, ok := m.topics[name]
+	m.topicsMux.RUnlock()
+
+	if !ok {
+		return nil, errors.Errorf("topic not found: %s", name)
+	}
+
+	ch := make(chan coretypes.ResultEvent)
+	m.subscribersMux.Lock()
+	defer m.subscribersMux.Unlock()
+	m.subscribers[name] = append(m.subscribers[name], ch)
+
+	return ch, nil
+}
+
+func (m *memEventBus) publishTopic(name string, src <-chan coretypes.ResultEvent) {
+	for {
+		select {
+		case msg, ok := <-src:
+			if !ok {
+				m.closeAllSubscribers(name)
+				m.topicsMux.Lock()
+				delete(m.topics, name)
+				m.topicsMux.Unlock()
+
+				return
+			}
+
+			m.publishAllSubscribers(name, msg)
+		}
+	}
+}
+
+func (m *memEventBus) closeAllSubscribers(name string) {
+	m.subscribersMux.Lock()
+	defer m.subscribersMux.Unlock()
+
+	subsribers := m.subscribers[name]
+	delete(m.subscribers, name)
+
+	for _, sub := range subsribers {
+		close(sub)
+	}
+}
+
+func (m *memEventBus) publishAllSubscribers(name string, msg coretypes.ResultEvent) {
+	m.subscribersMux.RLock()
+	subsribers := m.subscribers[name]
+	m.subscribersMux.RUnlock()
+
+	for _, sub := range subsribers {
+		select {
+		case sub <- msg:
+		default:
+		}
+	}
+}

ethereum/rpc/pubsub/pubsub_test.go

@@ -0,0 +1,96 @@
+package pubsub
+
+import (
+	"log"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestAddTopic(t *testing.T) {
+	assert := assert.New(t)
+
+	q := NewEventBus()
+	err := q.AddTopic("kek", make(chan interface{}))
+	if !assert.NoError(err) {
+		return
+	}
+
+	err = q.AddTopic("lol", make(chan interface{}))
+	if !assert.NoError(err) {
+		return
+	}
+
+	err = q.AddTopic("lol", make(chan interface{}))
+	if !assert.Error(err) {
+		return
+	}
+
+	assert.EqualValues([]string{"kek", "lol"}, q.Topics())
+}
+
+func TestSubscribe(t *testing.T) {
+	assert := assert.New(t)
+
+	q := NewEventBus()
+	kekSrc := make(chan interface{})
+	q.AddTopic("kek", kekSrc)
+
+	lolSrc := make(chan interface{})
+	q.AddTopic("lol", lolSrc)
+
+	kekSubC, err := q.Subscribe("kek")
+	if !assert.NoError(err) {
+		return
+	}
+
+	lolSubC, err := q.Subscribe("lol")
+	if !assert.NoError(err) {
+		return
+	}
+
+	lol2SubC, err := q.Subscribe("lol")
+	if !assert.NoError(err) {
+		return
+	}
+
+	wg := new(sync.WaitGroup)
+	wg.Add(4)
+
+	go func() {
+		defer wg.Done()
+		msg := <-kekSubC
+		log.Println("kek:", msg)
+		assert.EqualValues(1, msg)
+	}()
+
+	go func() {
+		defer wg.Done()
+		msg := <-lolSubC
+		log.Println("lol:", msg)
+		assert.EqualValues(1, msg)
+	}()
+
+	go func() {
+		defer wg.Done()
+		msg := <-lol2SubC
+		log.Println("lol2:", msg)
+		assert.EqualValues(1, msg)
+	}()
+
+	go func() {
+		defer wg.Done()
+
+		time.Sleep(time.Second)
+		kekSrc <- 1
+		lolSrc <- 1
+
+		close(kekSrc)
+		close(lolSrc)
+	}()
+
+	wg.Wait()
+	time.Sleep(time.Second)
+}