Add Update subscription type

Author: janos

update.go

@@ -0,0 +1,192 @@
+// Copyright (c) 2022, Janoš Guljaš <janos@resenje.org>
+// All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package feed
+
+import (
+	"sync"
+)
+
+// Update defines a set of subscriptions per topic T which receive messages sent
+// to the Update.
+type Update[T comparable, M any] struct {
+	subscriptions map[T][]*updateSubscription[T, M]
+	mu            sync.RWMutex
+
+	wg       sync.WaitGroup
+	quit     chan struct{}
+	quitOnce sync.Once
+}
+
+// NewFeed constructs new Feed with topic type T and message type M.
+func NewUpdate[T comparable, M any]() *Update[T, M] {
+	return &Update[T, M]{
+		subscriptions: make(map[T][]*updateSubscription[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 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 (u *Update[T, M]) Subscribe(topic T) (c <-chan M, cancel func()) {
+	channel := make(chan M)
+
+	select {
+	case <-u.quit:
+		close(channel)
+		return channel, func() {}
+	default:
+	}
+
+	u.mu.Lock()
+	defer u.mu.Unlock()
+
+	s := newUpdateSubscription(u, channel)
+
+	u.subscriptions[topic] = append(u.subscriptions[topic], s)
+
+	return channel, func() { u.unsubscribe(topic, s) }
+}
+
+func (u *Update[T, M]) unsubscribe(topic T, s *updateSubscription[T, M]) {
+	u.mu.Lock()
+	defer u.mu.Unlock()
+
+	for i, sub := range u.subscriptions[topic] {
+		if sub == s {
+			u.subscriptions[topic] = append(u.subscriptions[topic][:i], u.subscriptions[topic][i+1:]...)
+			s.close()
+		}
+	}
+}
+
+// Close terminates all subscriptions and releases acquired resources.
+func (u *Update[T, M]) Close() error {
+	u.quitOnce.Do(func() {
+		close(u.quit)
+	})
+	u.wg.Wait()
+
+	u.mu.Lock()
+	defer u.mu.Unlock()
+
+	for topic, subscriptions := range u.subscriptions {
+		for _, s := range subscriptions {
+			s.close()
+		}
+		u.subscriptions[topic] = nil
+	}
+
+	return nil
+}
+
+// Send sends a message to all sunscribed channels to topic. Messages will be
+// delivered to subscribers when each of them is ready to receive it, without
+// blocking this method call. The returned integer is the number of subscribers
+// that should receive the message.
+func (u *Update[T, M]) Send(topic T, message M) (n int) {
+	u.mu.RLock()
+	defer u.mu.RUnlock()
+
+	for _, s := range u.subscriptions[topic] {
+		s.send(message)
+
+		n++
+	}
+
+	return n
+}
+
+type updateSubscription[T comparable, M any] struct {
+	feed *Update[T, M]
+
+	channel chan M
+	update  chan M
+	updated chan struct{}
+
+	quit chan struct{}
+	wg   sync.WaitGroup
+}
+
+func newUpdateSubscription[T comparable, M any](u *Update[T, M], channel chan M) *updateSubscription[T, M] {
+	return &updateSubscription[T, M]{
+		feed:    u,
+		channel: channel,
+		update:  make(chan M),
+		updated: make(chan struct{}),
+		quit:    make(chan struct{}),
+	}
+}
+
+func (s *updateSubscription[T, M]) send(message M) {
+	select {
+	case s.channel <- message:
+
+	case s.update <- message:
+		select {
+		case <-s.updated:
+		case <-s.quit:
+		case <-s.feed.quit:
+		}
+		return
+
+	case <-s.quit:
+		return
+
+	case <-s.feed.quit:
+		return
+
+	default:
+
+		ready := make(chan struct{})
+		done := make(chan struct{})
+
+		channel := s.channel
+
+		s.wg.Add(1)
+		go func() {
+			defer s.wg.Done()
+			defer close(done)
+
+			for {
+				select {
+				case channel <- message:
+					return
+
+				case message = <-s.update:
+					channel = nil
+
+				case ready <- struct{}{}:
+
+				case s.updated <- struct{}{}:
+					channel = s.channel
+
+				case <-s.quit:
+					return
+
+				case <-s.feed.quit:
+					return
+
+				}
+			}
+		}()
+
+		select {
+		case <-ready:
+		case <-done:
+		case <-s.quit:
+		case <-s.feed.quit:
+		}
+	}
+}
+
+func (s *updateSubscription[T, M]) close() {
+	close(s.quit)
+	s.wg.Wait()
+	close(s.channel)
+}