Improve Feed with preserved ordering

Author: janos

feed.go

@@ -12,7 +12,7 @@ import (
 // Feed defines a set of subscriptions per topic T which receive messages sent
 // to the Feed.
 type Feed[T comparable, M any] struct {
-	subscriptions map[T][]*subscription[M]
+	subscriptions map[T][]*subscription[T, M]
 	mu            sync.RWMutex
 
 	wg       sync.WaitGroup
@@ -23,16 +23,16 @@ type Feed[T comparable, M any] struct {
 // NewFeed constructs new Feed with topic type T and message type M.
 func NewFeed[T comparable, M any]() *Feed[T, M] {
 	return &Feed[T, M]{
-		subscriptions: make(map[T][]*subscription[M]),
+		subscriptions: make(map[T][]*subscription[T, M]),
 		quit:          make(chan struct{}),
 	}
 }
 
 // Subscribe returns a channel from which messages M, that are sent to the Feed
-// on the same topic, can be read from. Message delivery is guaranteed, so the
-// channel should be read to avoid possible high number of goroutines. After
-// cancel function call, all resources ang goroutines are released even if not
-// all messages are read from channel.
+// on the same topic, can be read from. Message delivery preserves ordering and
+// is guaranteed, so the channel should be read to avoid keeping unread messages
+// in memory. After cancel function call, all resources ang goroutines are
+// released even if not all messages are read from channel.
 func (f *Feed[T, M]) Subscribe(topic T) (c <-chan M, cancel func()) {
 	channel := make(chan M, 1) // buffer 1 not to block on Send method
 
@@ -46,14 +46,14 @@ func (f *Feed[T, M]) Subscribe(topic T) (c <-chan M, cancel func()) {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 
-	s := newSubscription(channel)
+	s := newSubscription(f, channel)
 
 	f.subscriptions[topic] = append(f.subscriptions[topic], s)
 
 	return channel, func() { f.unsubscribe(topic, s) }
 }
 
-func (f *Feed[T, M]) unsubscribe(topic T, s *subscription[M]) {
+func (f *Feed[T, M]) unsubscribe(topic T, s *subscription[T, M]) {
 	f.mu.Lock()
 	defer f.mu.Unlock()
 
@@ -94,55 +94,85 @@ func (f *Feed[T, M]) Send(topic T, message M) (n int) {
 	defer f.mu.RUnlock()
 
 	for _, s := range f.subscriptions[topic] {
-		// try to send message to the channel
-		select {
-		case s.channel <- message:
-		case <-s.quit:
-			return
-		case <-f.quit:
-			return
-		default:
-			// if channel is blocked,
-			// wait in goroutine to send the message
-			s := s
-
-			f.wg.Add(1)
-			go func() {
-				defer f.wg.Done()
-
-				select {
-				case s.channel <- message:
-				case <-s.quit:
-					return
-				case <-f.quit:
-					return
-				}
-			}()
-		}
+		s.send(message)
 
 		n++
 	}
 
 	return n
 }
 
-type subscription[M any] struct {
-	channel  chan M
-	quitOnce sync.Once
-	quit     chan struct{}
-	wg       sync.WaitGroup
+type subscription[T comparable, M any] struct {
+	feed *Feed[T, M]
+
+	channel chan M
+
+	buf []M
+	mu  sync.Mutex
+
+	quit chan struct{}
+	wg   sync.WaitGroup
 }
 
-func newSubscription[M any](channel chan M) *subscription[M] {
-	return &subscription[M]{
+func newSubscription[T comparable, M any](feed *Feed[T, M], channel chan M) *subscription[T, M] {
+	return &subscription[T, M]{
+		feed:    feed,
 		channel: channel,
+		buf:     make([]M, 0),
 		quit:    make(chan struct{}),
 	}
 }
 
-func (s *subscription[M]) close() {
-	s.quitOnce.Do(func() {
-		close(s.quit)
-	})
+func (s *subscription[T, M]) send(message M) {
+	s.mu.Lock()
+	defer s.mu.Unlock()
+
+	if len(s.buf) > 0 {
+		s.buf = append(s.buf, message)
+		return
+	}
+
+	select {
+	case s.channel <- message:
+	case <-s.quit:
+		return
+	case <-s.feed.quit:
+		return
+	default:
+		s.buf = append(s.buf, message)
+
+		s.wg.Add(1)
+		go func() {
+			defer s.wg.Done()
+
+			for {
+				s.mu.Lock()
+				message := s.buf[0]
+				s.mu.Unlock()
+
+				select {
+				case s.channel <- message:
+				case <-s.quit:
+					return
+				case <-s.feed.quit:
+					return
+				}
+
+				s.mu.Lock()
+				s.buf = s.buf[1:]
+				if len(s.buf) == 0 {
+					s.mu.Unlock()
+					return
+				}
+				s.mu.Unlock()
+			}
+
+		}()
+	}
+}
+
+func (s *subscription[T, M]) close() {
+	close(s.quit)
 	s.wg.Wait()
+	close(s.channel)
 }

feed_test.go

@@ -6,9 +6,11 @@
 package feed_test
 
 import (
+	"math/rand"
 	"reflect"
 	"sync"
 	"testing"
+	"time"
 
 	"resenje.org/feed"
 )
@@ -263,6 +265,78 @@ func TestFeed_cancelWithUnreadMessages(t *testing.T) {
 	cancel()
 }
 
+func TestFeed_ordering(t *testing.T) {
+	f := feed.NewFeed[string, int]()
+	defer f.Close()
+
+	got := make([]int, 0)
+
+	s, cancel := f.Subscribe("topic1")
+	defer cancel()
+
+	messages := make([]int, 0)
+
+	count := 1000
+
+	for i := 0; i < count; i++ {
+		messages = append(messages, i)
+		n := f.Send("topic1", i)
+		assert(t, "", n, 1)
+	}
+
+	for i := 0; i < count; i++ {
+		got = append(got, <-s)
+	}
+
+	assert(t, "", got, messages)
+}
+
+func TestFeed_stressTest(t *testing.T) {
+	f := feed.NewFeed[string, int]()
+	defer f.Close()
+
+	count := 1000
+	lastReceived := -1
+
+	s, cancel := f.Subscribe("topic1")
+	defer cancel()
+
+	random1 := rand.New(rand.NewSource(time.Now().UnixNano()))
+	random2 := rand.New(rand.NewSource(time.Now().UnixNano()))
+
+	receiveDone := make(chan struct{})
+	go func() {
+		defer close(receiveDone)
+		for {
+			m, ok := <-s
+			if !ok {
+				return
+			}
+			if m != lastReceived+1 {
+				assert(t, "", m, lastReceived+1)
+			}
+			lastReceived = m
+			time.Sleep(time.Duration(random1.Int63n(1000)) * time.Microsecond)
+		}
+	}()
+
+	sendDone := make(chan struct{})
+	go func() {
+		defer close(sendDone)
+		for i := 0; i < count; i++ {
+			n := f.Send("topic1", i)
+			assert(t, "", n, 1)
+			time.Sleep(time.Duration(random2.Int63n(1000)) * time.Microsecond)
+		}
+	}()
+
+	<-sendDone
+
+	cancel()
+
+	<-receiveDone
+}
+
 func assert[T any](t testing.TB, message string, got, want T) {
 	t.Helper()