Multi-watcher with round robin orchestrator

hub/hub.go

@@ -12,84 +12,108 @@ import (
 // Hub - Abstraction between message publishers & subscribers,
 // works as a multiplexer ( or router )
 type Hub struct {
-	addr                       string
-	watcher                    *gaio.Watcher
-	pendingPublishers          map[net.Conn]bool
-	pendingNewSubscribers      map[net.Conn]bool
-	pendingExistingSubscribers map[net.Conn]bool
-	pendingUnsubscribers       map[net.Conn]bool
-	enqueuedReadLock           *sync.RWMutex
-	enqueuedRead               map[net.Conn]*enqueuedRead
-	connectedSubscribers       map[net.Conn]uint64
-	index                      uint64
-	subLock                    *sync.RWMutex
-	subscribers                map[string]map[uint64]net.Conn
-	revLock                    *sync.RWMutex
-	revSubscribers             map[uint64]map[string]bool
-	queueLock                  *sync.RWMutex
-	pendingQueue               []*ops.Msg
-	ping                       chan struct{}
-	evict                      chan uint64
-	Connected                  chan string
-	Disconnected               chan string
+	addr                     string
+	watchersLock             *sync.RWMutex
+	watchers                 map[uint]*watcher
+	watcherCount             uint
+	connectedSubscribers     map[net.Conn]*subInfo
+	connectedSubscribersLock *sync.RWMutex
+	index                    uint64
+	subLock                  *sync.RWMutex
+	subscribers              map[string]map[uint64]net.Conn
+	revLock                  *sync.RWMutex
+	revSubscribers           map[uint64]map[string]bool
+	queueLock                *sync.RWMutex
+	pendingQueue             []*ops.Msg
+	ping                     chan struct{}
+	evict                    chan uint64
+	Connected                chan string
+	Disconnected             chan string
 }
 
 func (h *Hub) Addr() string {
 	return h.addr
 }
 
-type enqueuedRead struct {
-	yes bool
-	buf []byte
+type watcher struct {
+	eventLoop   *gaio.Watcher
+	ongoingRead map[net.Conn]*readState
+	lock        *sync.RWMutex
+}
+
+type readState struct {
+	opcode       ops.OP
+	envelopeRead bool
+	buf          []byte
+}
+
+type subInfo struct {
+	id        uint64
+	watcherId uint
 }
 
 // New - Creates a new instance of hub, ready to be used
 func New(ctx context.Context, addr string, cap uint64) (*Hub, error) {
-	watcher, err := gaio.NewWatcher()
-	if err != nil {
-		return nil, err
+	hub := Hub{
+		watcherCount:             2,
+		watchersLock:             &sync.RWMutex{},
+		watchers:                 make(map[uint]*watcher),
+		connectedSubscribers:     make(map[net.Conn]*subInfo),
+		connectedSubscribersLock: &sync.RWMutex{},
+		index:                    0,
+		subLock:                  &sync.RWMutex{},
+		subscribers:              make(map[string]map[uint64]net.Conn),
+		revLock:                  &sync.RWMutex{},
+		revSubscribers:           make(map[uint64]map[string]bool),
+		queueLock:                &sync.RWMutex{},
+		pendingQueue:             make([]*ops.Msg, 0, cap),
+		ping:                     make(chan struct{}, cap),
+		evict:                    make(chan uint64, cap),
+		Connected:                make(chan string, 1),
+		Disconnected:             make(chan string, 1),
 	}
 
-	hub := Hub{
-		watcher:                    watcher,
-		pendingPublishers:          make(map[net.Conn]bool),
-		pendingNewSubscribers:      make(map[net.Conn]bool),
-		pendingExistingSubscribers: make(map[net.Conn]bool),
-		pendingUnsubscribers:       make(map[net.Conn]bool),
-		enqueuedReadLock:           &sync.RWMutex{},
-		enqueuedRead:               make(map[net.Conn]*enqueuedRead),
-		connectedSubscribers:       make(map[net.Conn]uint64),
-		index:                      0,
-		subLock:                    &sync.RWMutex{},
-		subscribers:                make(map[string]map[uint64]net.Conn),
-		revLock:                    &sync.RWMutex{},
-		revSubscribers:             make(map[uint64]map[string]bool),
-		queueLock:                  &sync.RWMutex{},
-		pendingQueue:               make([]*ops.Msg, 0, cap),
-		ping:                       make(chan struct{}, cap),
-		evict:                      make(chan uint64, cap),
-		Connected:                  make(chan string, 1),
-		Disconnected:               make(chan string, 1),
+	var runWatcher = make(chan struct{}, hub.watcherCount)
+	var i uint
+	hub.watchersLock.Lock()
+	for ; i < hub.watcherCount; i++ {
+		w, err := gaio.NewWatcher()
+		if err != nil {
+			return nil, err
+		}
+		hub.watchers[i] = &watcher{
+			eventLoop:   w,
+			ongoingRead: make(map[net.Conn]*readState),
+			lock:        &sync.RWMutex{},
+		}
+		func(id uint) {
+			go hub.watch(ctx, id, runWatcher)
+		}(i)
+	}
+	hub.watchersLock.Unlock()
+
+	startedOff := 0
+	for range runWatcher {
+		startedOff++
+		if startedOff >= int(hub.watcherCount) {
+			break
+		}
 	}
 
 	var (
-		runListener = make(chan bool)
-		runWatcher  = make(chan struct{})
-		runProc     = make(chan struct{})
-		runEvict    = make(chan struct{})
+		runListener  = make(chan bool)
+		runEvictor   = make(chan struct{})
+		runProcessor = make(chan struct{})
 	)
 
 	go hub.listen(ctx, addr, runListener)
 	if !<-runListener {
 		return nil, ops.ErrListenerNotStarted
 	}
-
-	go hub.watch(ctx, runWatcher)
-	go hub.process(ctx, runProc)
-	go hub.evictSubscribers(ctx, runEvict)
-	<-runWatcher
-	<-runProc
-	<-runEvict
+	go hub.process(ctx, runProcessor)
+	go hub.evictSubscribers(ctx, runEvictor)
+	<-runProcessor
+	<-runEvictor
 
 	return &hub, nil
 }

hub/listener.go

@@ -5,6 +5,8 @@ import (
 	"fmt"
 	"log"
 	"net"
+
+	"github.com/itzmeanjan/pub0sub/ops"
 )
 
 func (h *Hub) listen(ctx context.Context, addr string, done chan bool) {
@@ -24,6 +26,7 @@ func (h *Hub) listen(ctx context.Context, addr string, done chan bool) {
 
 	h.addr = lis.Addr().String()
 	done <- true
+	var nextWatcher uint
 
 	for {
 		select {
@@ -42,12 +45,18 @@ func (h *Hub) listen(ctx context.Context, addr string, done chan bool) {
 				h.Connected <- fmt.Sprintf("%s://%s", conn.RemoteAddr().Network(), conn.RemoteAddr().String())
 			}
 
-			h.enqueuedReadLock.Lock()
 			buf := make([]byte, 5)
-			h.enqueuedRead[conn] = &enqueuedRead{yes: true, buf: buf}
-			h.enqueuedReadLock.Unlock()
+			nextWatcher = (nextWatcher + 1) % h.watcherCount
+
+			h.watchersLock.RLock()
+			watcher := h.watchers[nextWatcher]
+			h.watchersLock.RUnlock()
+
+			watcher.lock.Lock()
+			watcher.ongoingRead[conn] = &readState{buf: buf, opcode: ops.UNSUPPORTED}
+			watcher.lock.Unlock()
 
-			if err := h.watcher.Read(ctx, conn, buf); err != nil {
+			if err := watcher.eventLoop.Read(ctx, conn, buf); err != nil {
 				log.Printf("[pub0sub] Error : %s\n", err.Error())
 				return
 			}

hub/processor.go

@@ -63,10 +63,15 @@ func (h *Hub) writeMessage(ctx context.Context, op *ops.OP, msg *ops.Msg) {
 			continue
 		}
 
+		h.connectedSubscribersLock.RLock()
+		defer h.connectedSubscribersLock.RUnlock()
+
 		for _, conn := range subs {
-			if err := h.watcher.Write(ctx, conn, buf.Bytes()); err != nil {
-				log.Printf("[pub0sub] Error : %s\n", err.Error())
-				continue
+			if subInfo, ok := h.connectedSubscribers[conn]; ok {
+				if err := h.watchers[subInfo.watcherId].eventLoop.Write(ctx, conn, buf.Bytes()); err != nil {
+					log.Printf("[pub0sub] Error : %s\n", err.Error())
+					continue
+				}
 			}
 		}
 	}