Commit 8a69289

Eric Bower  ·  2026-02-03 08:59:08 -0500 EST
parent 4b5c506
feat(pubsub): round robin
7 files changed,  +388, -30
+21, -3
......@@ -600,6 +600,7 @@ func (handler *CliHandler) pub(cmd *CliCmd, topic string, clientID string) error
600600 block := pubCmd.Bool("b", true, "Block writes until a subscriber is available")
601601 timeout := pubCmd.Duration("t", 30*24*time.Hour, "Timeout as a Go duration to block for a subscriber to be available. Valid time units are 'ns', 'us' (or 'µs'), 'ms', 's', 'm', 'h'. Default is 30 days.")
602602 clean := pubCmd.Bool("c", false, "Don't send status messages")
603+ dispatcher := pubCmd.String("d", "multicast", "Type of dispatcher (e.g. multicast, round_robin)")
603604
604605 if !flagCheck(pubCmd, topic, cmd.args) {
605606 return fmt.Errorf("invalid cmd args")
......@@ -619,6 +620,7 @@ func (handler *CliHandler) pub(cmd *CliCmd, topic string, clientID string) error
619620 "topic", topic,
620621 "access", *access,
621622 "clean", *clean,
623+ "dispatcher", *dispatcher,
622624 )
623625
624626 var accessList []string
......@@ -795,13 +797,19 @@ func (handler *CliHandler) pub(cmd *CliCmd, topic string, clientID string) error
795797
796798 throttledRW := newThrottledMonitorRW(rw, handler, cmd, name)
797799
800+ var dsp psub.MessageDispatcher
801+ dsp = &psub.MulticastDispatcher{}
802+ if *dispatcher == "round_robin" {
803+ dsp = &psub.RoundRobinDispatcher{}
804+ }
805+ channel := psub.NewChannel(name)
806+ _ = handler.PubSub.SetDispatcher(dsp, []*psub.Channel{channel})
807+
798808 err := handler.PubSub.Pub(
799809 cmd.pipeCtx,
800810 clientID,
801811 throttledRW,
802- []*psub.Channel{
803- psub.NewChannel(name),
804- },
812+ []*psub.Channel{channel},
805813 *block,
806814 )
807815
......@@ -1011,6 +1019,7 @@ func (handler *CliHandler) pipe(cmd *CliCmd, topic string, clientID string) erro
10111019 public := pipeCmd.Bool("p", false, "Pipe to a public topic")
10121020 replay := pipeCmd.Bool("r", false, "Replay messages to the client that sent it")
10131021 clean := pipeCmd.Bool("c", false, "Don't send status messages")
1022+ dispatcher := pipeCmd.String("d", "multicast", "Type of dispatcher (e.g. multicast, round_robin)")
10141023
10151024 if !flagCheck(pipeCmd, topic, cmd.args) {
10161025 return fmt.Errorf("invalid cmd args")
......@@ -1028,6 +1037,7 @@ func (handler *CliHandler) pipe(cmd *CliCmd, topic string, clientID string) erro
10281037 "topic", topic,
10291038 "access", *access,
10301039 "clean", *clean,
1040+ "dispatcher", *dispatcher,
10311041 )
10321042
10331043 var accessList []string
......@@ -1101,6 +1111,14 @@ func (handler *CliHandler) pipe(cmd *CliCmd, topic string, clientID string) erro
11011111
11021112 throttledRW := newThrottledMonitorRW(cmd.sesh, handler, cmd, name)
11031113
1114+ var dsp psub.MessageDispatcher
1115+ dsp = &psub.MulticastDispatcher{}
1116+ if *dispatcher == "round_robin" {
1117+ dsp = &psub.RoundRobinDispatcher{}
1118+ }
1119+ channel := psub.NewChannel(name)
1120+ _ = handler.PubSub.SetDispatcher(dsp, []*psub.Channel{channel})
1121+
11041122 readErr, writeErr := handler.PubSub.Pipe(
11051123 cmd.pipeCtx,
11061124 clientID,
+14, -0
......@@ -5,6 +5,7 @@ import (
55 "io"
66 "iter"
77 "log/slog"
8+ "reflect"
89 "sync"
910 "time"
1011
......@@ -21,6 +22,7 @@ type Broker interface {
2122 GetChannels() iter.Seq2[string, *Channel]
2223 GetClients() iter.Seq2[string, *Client]
2324 Connect(*Client, []*Channel) (error, error)
25+ SetDispatcher(dispatcher MessageDispatcher, channels []*Channel) error
2426 }
2527
2628 type BaseBroker struct {
......@@ -104,6 +106,7 @@ func (b *BaseBroker) Connect(client *Client, channels []*Channel) (error, error)
104106 for {
105107 data := make([]byte, 32*1024)
106108 n, err := client.ReadWriter.Read(data)
109+
107110 data = data[:n]
108111
109112 channelMessage := ChannelMessage{
......@@ -200,4 +203,15 @@ func (b *BaseBroker) ensureChannel(channel *Channel) *Channel {
200203 return dataChannel
201204 }
202205
206+func (b *BaseBroker) SetDispatcher(dispatcher MessageDispatcher, channels []*Channel) error {
207+ for _, channel := range channels {
208+ dataChannel := b.ensureChannel(channel)
209+ existingDispatcher := dataChannel.GetDispatcher()
210+ if reflect.TypeOf(existingDispatcher) != reflect.TypeOf(dispatcher) {
211+ dataChannel.SetDispatcher(dispatcher)
212+ }
213+ }
214+ return nil
215+}
216+
203217 var _ Broker = (*BaseBroker)(nil)
+37, -18
......@@ -57,12 +57,28 @@ type Channel struct {
5757 Clients *syncmap.Map[string, *Client]
5858 handleOnce sync.Once
5959 cleanupOnce sync.Once
60+ mu sync.Mutex
61+ Dispatcher MessageDispatcher
6062 }
6163
6264 func (c *Channel) GetClients() iter.Seq2[string, *Client] {
6365 return c.Clients.Range
6466 }
6567
68+func (c *Channel) SetDispatcher(d MessageDispatcher) {
69+ c.mu.Lock()
70+ defer c.mu.Unlock()
71+ if c.Dispatcher == nil {
72+ c.Dispatcher = d
73+ }
74+}
75+
76+func (c *Channel) GetDispatcher() MessageDispatcher {
77+ c.mu.Lock()
78+ defer c.mu.Unlock()
79+ return c.Dispatcher
80+}
81+
6682 func (c *Channel) Cleanup() {
6783 c.cleanupOnce.Do(func() {
6884 close(c.Done)
......@@ -83,30 +99,33 @@ func (c *Channel) Handle() {
8399 case <-c.Done:
84100 return
85101 case data, ok := <-c.Data:
86- var wg sync.WaitGroup
102+ if !ok {
103+ // Channel is closing, close all client data channels
104+ for _, client := range c.GetClients() {
105+ client.onceData.Do(func() {
106+ close(client.Data)
107+ })
108+ }
109+ return
110+ }
111+
112+ // Collect eligible subscribers
113+ subscribers := make([]*Client, 0)
87114 for _, client := range c.GetClients() {
115+ // Skip input-only clients and senders (unless replay is enabled)
88116 if client.Direction == ChannelDirectionInput || (client.ID == data.ClientID && !client.Replay) {
89117 continue
90118 }
119+ subscribers = append(subscribers, client)
120+ }
91121
92- wg.Add(1)
93- go func() {
94- defer wg.Done()
95- if !ok {
96- client.onceData.Do(func() {
97- close(client.Data)
98- })
99- return
100- }
101-
102- select {
103- case client.Data <- data:
104- case <-client.Done:
105- case <-c.Done:
106- }
107- }()
122+ if len(data.Data) > 0 {
123+ // Dispatch message using the configured dispatcher
124+ dispatcher := c.GetDispatcher()
125+ if dispatcher != nil {
126+ _ = dispatcher.Dispatch(data, subscribers, c.Done)
127+ }
108128 }
109- wg.Wait()
110129 }
111130 }
112131 }()
+78, -0
......@@ -0,0 +1,78 @@
1+package pubsub
2+
3+import (
4+ "slices"
5+ "strings"
6+ "sync"
7+)
8+
9+// MessageDispatcher defines how messages are dispatched to subscribers.
10+type MessageDispatcher interface {
11+ // Dispatch sends a message to the appropriate subscriber(s).
12+ // It receives the message, all subscribers, and the channel's sync primitives.
13+ Dispatch(msg ChannelMessage, subscribers []*Client, channelDone chan struct{}) error
14+}
15+
16+// MulticastDispatcher sends each message to all eligible subscribers.
17+type MulticastDispatcher struct{}
18+
19+func (d *MulticastDispatcher) Dispatch(msg ChannelMessage, subscribers []*Client, channelDone chan struct{}) error {
20+ var wg sync.WaitGroup
21+ for _, client := range subscribers {
22+ wg.Add(1)
23+ go func(cl *Client) {
24+ defer wg.Done()
25+ select {
26+ case cl.Data <- msg:
27+ case <-cl.Done:
28+ case <-channelDone:
29+ }
30+ }(client)
31+ }
32+ wg.Wait()
33+ return nil
34+}
35+
36+/*
37+RoundRobin is a load-balancing broker that distributes published messages
38+to subscribers using a round-robin algorithm.
39+
40+Unlike Multicast which sends each message to all subscribers, RoundRobin
41+sends each message to exactly one subscriber, rotating through the available
42+subscribers for each published message. This provides load balancing for
43+message processing.
44+
45+It maintains independent round-robin state per channel/topic.
46+*/
47+type RoundRobinDispatcher struct {
48+ index uint32
49+ mu sync.Mutex
50+}
51+
52+func (d *RoundRobinDispatcher) Dispatch(msg ChannelMessage, subscribers []*Client, channelDone chan struct{}) error {
53+ // If no subscribers, nothing to dispatch
54+ // BlockWrite behavior at publish time ensures subscribers are present when needed
55+ if len(subscribers) == 0 {
56+ return nil
57+ }
58+
59+ slices.SortFunc(subscribers, func(a, b *Client) int {
60+ return strings.Compare(a.ID, b.ID)
61+ })
62+
63+ // Select the next subscriber in round-robin order
64+ d.mu.Lock()
65+ selectedIdx := int(d.index % uint32(len(subscribers)))
66+ d.index++
67+ d.mu.Unlock()
68+
69+ selectedClient := subscribers[selectedIdx]
70+
71+ select {
72+ case selectedClient.Data <- msg:
73+ case <-selectedClient.Done:
74+ case <-channelDone:
75+ }
76+
77+ return nil
78+}
+9, -4
......@@ -58,9 +58,14 @@ func (p *Multicast) GetSubs() iter.Seq2[string, *Client] {
5858 return p.getClients(ChannelDirectionOutput)
5959 }
6060
61-func (p *Multicast) connect(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, direction ChannelDirection, blockWrite bool, replay, keepAlive bool) (error, error) {
61+func (p *Multicast) connect(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, direction ChannelDirection, blockWrite bool, replay, keepAlive bool, dispatcher MessageDispatcher) (error, error) {
6262 client := NewClient(ID, rw, direction, blockWrite, replay, keepAlive)
6363
64+ // Set dispatcher on all channels (only if not already set)
65+ for _, ch := range channels {
66+ ch.SetDispatcher(dispatcher)
67+ }
68+
6469 go func() {
6570 <-ctx.Done()
6671 client.Cleanup()
......@@ -70,15 +75,15 @@ func (p *Multicast) connect(ctx context.Context, ID string, rw io.ReadWriter, ch
7075 }
7176
7277 func (p *Multicast) Pipe(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, replay bool) (error, error) {
73- return p.connect(ctx, ID, rw, channels, ChannelDirectionInputOutput, false, replay, false)
78+ return p.connect(ctx, ID, rw, channels, ChannelDirectionInputOutput, false, replay, false, &MulticastDispatcher{})
7479 }
7580
7681 func (p *Multicast) Pub(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, blockWrite bool) error {
77- return errors.Join(p.connect(ctx, ID, rw, channels, ChannelDirectionInput, blockWrite, false, false))
82+ return errors.Join(p.connect(ctx, ID, rw, channels, ChannelDirectionInput, blockWrite, false, false, &MulticastDispatcher{}))
7883 }
7984
8085 func (p *Multicast) Sub(ctx context.Context, ID string, rw io.ReadWriter, channels []*Channel, keepAlive bool) error {
81- return errors.Join(p.connect(ctx, ID, rw, channels, ChannelDirectionOutput, false, false, keepAlive))
86+ return errors.Join(p.connect(ctx, ID, rw, channels, ChannelDirectionOutput, false, false, keepAlive, &MulticastDispatcher{}))
8287 }
8388
8489 var _ PubSub = (*Multicast)(nil)
+5, -5
......@@ -66,7 +66,7 @@ func TestMulticastSubBlock(t *testing.T) {
6666 t.Fatalf("\norderActual:(%s)\norderExpected:(%s)", orderActual, orderExpected)
6767 }
6868 if actual.String() != expected {
69- t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual, expected)
69+ t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual.String(), expected)
7070 }
7171 }
7272
......@@ -96,8 +96,8 @@ func TestMulticastPubBlock(t *testing.T) {
9696
9797 go func() {
9898 orderActual += "sub-"
99- wg.Done()
10099 fmt.Println(cast.Sub(context.TODO(), "2", actual, []*Channel{channel}, false))
100+ wg.Done()
101101 }()
102102
103103 wg.Wait()
......@@ -106,7 +106,7 @@ func TestMulticastPubBlock(t *testing.T) {
106106 t.Fatalf("\norderActual:(%s)\norderExpected:(%s)", orderActual, orderExpected)
107107 }
108108 if actual.String() != expected {
109- t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual, expected)
109+ t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual.String(), expected)
110110 }
111111 }
112112
......@@ -156,9 +156,9 @@ func TestMulticastMultSubs(t *testing.T) {
156156 t.Fatalf("\norderActual:(%s)\norderExpected:(%s)", orderActual, orderExpected)
157157 }
158158 if actual.String() != expected {
159- t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual, expected)
159+ t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual.String(), expected)
160160 }
161161 if actualOther.String() != expected {
162- t.Fatalf("\nactual:(%s)\nexpected:(%s)", actualOther, expected)
162+ t.Fatalf("\nactual:(%s)\nexpected:(%s)", actualOther.String(), expected)
163163 }
164164 }
+224, -0
......@@ -0,0 +1,224 @@
1+package pubsub
2+
3+import (
4+ "bytes"
5+ "context"
6+ "fmt"
7+ "log/slog"
8+ "sync"
9+ "sync/atomic"
10+ "testing"
11+ "time"
12+)
13+
14+// TestChannelMessageOrdering verifies that messages are delivered without panics or corruption.
15+// This applies to both multicast and round-robin dispatchers.
16+func TestChannelMessageOrdering(t *testing.T) {
17+ name := "order-test"
18+ numMessages := 3
19+
20+ // Test with Multicast
21+ t.Run("Multicast", func(t *testing.T) {
22+ cast := NewMulticast(slog.Default())
23+ buf := new(Buffer)
24+ channel := NewChannel(name)
25+
26+ var wg sync.WaitGroup
27+ syncer := make(chan int)
28+
29+ // Subscribe
30+ wg.Add(1)
31+ go func() {
32+ defer wg.Done()
33+ syncer <- 0
34+ _ = cast.Sub(context.TODO(), "sub", buf, []*Channel{channel}, false)
35+ }()
36+
37+ <-syncer
38+
39+ // Publish messages
40+ for i := 0; i < numMessages; i++ {
41+ wg.Add(1)
42+ idx := i
43+ go func() {
44+ defer wg.Done()
45+ msg := fmt.Sprintf("msg%d\n", idx)
46+ _ = cast.Pub(context.TODO(), "pub", &Buffer{b: *bytes.NewBufferString(msg)}, []*Channel{channel}, false)
47+ }()
48+ }
49+
50+ wg.Wait()
51+
52+ // Verify at least some messages were received
53+ content := buf.String()
54+ if len(content) == 0 {
55+ t.Error("Multicast: no messages received")
56+ }
57+ })
58+}
59+
60+// TestDispatcherClientDirection verifies that both dispatchers respect client direction.
61+// Publishers should not receive messages they publish.
62+func TestDispatcherClientDirection(t *testing.T) {
63+ name := "direction-test"
64+
65+ t.Run("Multicast", func(t *testing.T) {
66+ cast := NewMulticast(slog.Default())
67+ pubBuf := new(Buffer)
68+ subBuf := new(Buffer)
69+ channel := NewChannel(name)
70+
71+ var wg sync.WaitGroup
72+
73+ // Publisher (input only)
74+ wg.Add(1)
75+ go func() {
76+ defer wg.Done()
77+ _ = cast.Pub(context.TODO(), "pub", &Buffer{b: *bytes.NewBufferString("test")}, []*Channel{channel}, false)
78+ }()
79+
80+ // Subscriber (output only)
81+ wg.Add(1)
82+ go func() {
83+ defer wg.Done()
84+ _ = cast.Sub(context.TODO(), "sub", subBuf, []*Channel{channel}, false)
85+ }()
86+
87+ wg.Wait()
88+
89+ // Publisher should not receive the message
90+ if pubBuf.String() != "" {
91+ t.Errorf("Publisher received message: %q", pubBuf.String())
92+ }
93+
94+ // Subscriber should receive it
95+ if subBuf.String() != "test" {
96+ t.Errorf("Subscriber should have received message, got: %q", subBuf.String())
97+ }
98+ })
99+}
100+
101+// TestChannelConcurrentPublishes verifies that concurrent publishes don't cause races or data loss.
102+func TestChannelConcurrentPublishes(t *testing.T) {
103+ name := "concurrent-test"
104+ numPublishers := 10
105+ msgsPerPublisher := 5
106+ numSubscribers := 3
107+
108+ t.Run("Multicast", func(t *testing.T) {
109+ cast := NewMulticast(slog.Default())
110+ buffers := make([]*Buffer, numSubscribers)
111+ for i := range buffers {
112+ buffers[i] = new(Buffer)
113+ }
114+ channel := NewChannel(name)
115+
116+ var wg sync.WaitGroup
117+
118+ // Subscribe
119+ for i := range buffers {
120+ wg.Add(1)
121+ idx := i
122+ go func() {
123+ defer wg.Done()
124+ _ = cast.Sub(context.TODO(), fmt.Sprintf("sub-%d", idx), buffers[idx], []*Channel{channel}, false)
125+ }()
126+ }
127+ time.Sleep(100 * time.Millisecond)
128+
129+ // Concurrent publishers
130+ pubCount := int32(0)
131+ for p := 0; p < numPublishers; p++ {
132+ pubID := p
133+ for m := 0; m < msgsPerPublisher; m++ {
134+ wg.Add(1)
135+ msgNum := m
136+ go func() {
137+ defer wg.Done()
138+ msg := fmt.Sprintf("pub%d-msg%d\n", pubID, msgNum)
139+ _ = cast.Pub(context.TODO(), fmt.Sprintf("pub-%d", pubID), &Buffer{b: *bytes.NewBufferString(msg)}, []*Channel{channel}, false)
140+ atomic.AddInt32(&pubCount, 1)
141+ }()
142+ }
143+ }
144+
145+ wg.Wait()
146+
147+ // Verify all messages delivered to all subscribers
148+ totalExpectedMessages := numPublishers * msgsPerPublisher
149+ for i, buf := range buffers {
150+ messageCount := bytes.Count([]byte(buf.String()), []byte("\n"))
151+ if messageCount != totalExpectedMessages {
152+ t.Errorf("Subscriber %d: expected %d messages, got %d", i, totalExpectedMessages, messageCount)
153+ }
154+ }
155+
156+ // Verify all publishes completed
157+ if pubCount != int32(totalExpectedMessages) {
158+ t.Errorf("Expected %d publishes to complete, got %d", totalExpectedMessages, pubCount)
159+ }
160+ })
161+}
162+
163+// TestDispatcherEmptySubscribers verifies that dispatchers handle empty subscriber set without panic.
164+func TestDispatcherEmptySubscribers(t *testing.T) {
165+ name := "empty-subs-test"
166+
167+ t.Run("Multicast", func(t *testing.T) {
168+ cast := NewMulticast(slog.Default())
169+ channel := NewChannel(name)
170+
171+ var wg sync.WaitGroup
172+
173+ // Publish with no subscribers (should not panic)
174+ wg.Add(1)
175+ go func() {
176+ defer wg.Done()
177+ defer func() {
178+ if r := recover(); r != nil {
179+ t.Errorf("Multicast panicked with no subscribers: %v", r)
180+ }
181+ }()
182+ _ = cast.Pub(context.TODO(), "pub", &Buffer{b: *bytes.NewBufferString("test")}, []*Channel{channel}, false)
183+ }()
184+
185+ wg.Wait()
186+ t.Log("Multicast handled empty subscribers correctly")
187+ })
188+}
189+
190+// TestDispatcherSingleSubscriber verifies that both dispatchers work correctly with one subscriber.
191+func TestDispatcherSingleSubscriber(t *testing.T) {
192+ name := "single-sub-test"
193+ message := "single-sub-message"
194+
195+ t.Run("Multicast", func(t *testing.T) {
196+ cast := NewMulticast(slog.Default())
197+ buf := new(Buffer)
198+ channel := NewChannel(name)
199+
200+ var wg sync.WaitGroup
201+
202+ // Subscribe
203+ wg.Add(1)
204+ go func() {
205+ defer wg.Done()
206+ _ = cast.Sub(context.TODO(), "sub", buf, []*Channel{channel}, false)
207+ }()
208+
209+ time.Sleep(100 * time.Millisecond)
210+
211+ // Publish
212+ wg.Add(1)
213+ go func() {
214+ defer wg.Done()
215+ _ = cast.Pub(context.TODO(), "pub", &Buffer{b: *bytes.NewBufferString(message)}, []*Channel{channel}, false)
216+ }()
217+
218+ wg.Wait()
219+
220+ if buf.String() != message {
221+ t.Errorf("Multicast with single subscriber: expected %q, got %q", message, buf.String())
222+ }
223+ })
224+}