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 | |
| 600 | 600 | block := pubCmd.Bool("b", true, "Block writes until a subscriber is available") | |
| 601 | 601 | 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.") | |
| 602 | 602 | clean := pubCmd.Bool("c", false, "Don't send status messages") | |
| 603 | + | dispatcher := pubCmd.String("d", "multicast", "Type of dispatcher (e.g. multicast, round_robin)") | |
| 603 | 604 | ||
| 604 | 605 | if !flagCheck(pubCmd, topic, cmd.args) { | |
| 605 | 606 | return fmt.Errorf("invalid cmd args") |
| ... | ... | @@ -619,6 +620,7 @@ func (handler *CliHandler) pub(cmd *CliCmd, topic string, clientID string) error | |
| 619 | 620 | "topic", topic, | |
| 620 | 621 | "access", *access, | |
| 621 | 622 | "clean", *clean, | |
| 623 | + | "dispatcher", *dispatcher, | |
| 622 | 624 | ) | |
| 623 | 625 | ||
| 624 | 626 | var accessList []string |
| ... | ... | @@ -795,13 +797,19 @@ func (handler *CliHandler) pub(cmd *CliCmd, topic string, clientID string) error | |
| 795 | 797 | ||
| 796 | 798 | throttledRW := newThrottledMonitorRW(rw, handler, cmd, name) | |
| 797 | 799 | ||
| 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 | + | ||
| 798 | 808 | err := handler.PubSub.Pub( | |
| 799 | 809 | cmd.pipeCtx, | |
| 800 | 810 | clientID, | |
| 801 | 811 | throttledRW, | |
| 802 | - | []*psub.Channel{ | |
| 803 | - | psub.NewChannel(name), | |
| 804 | - | }, | |
| 812 | + | []*psub.Channel{channel}, | |
| 805 | 813 | *block, | |
| 806 | 814 | ) | |
| 807 | 815 |
| ... | ... | @@ -1011,6 +1019,7 @@ func (handler *CliHandler) pipe(cmd *CliCmd, topic string, clientID string) erro | |
| 1011 | 1019 | public := pipeCmd.Bool("p", false, "Pipe to a public topic") | |
| 1012 | 1020 | replay := pipeCmd.Bool("r", false, "Replay messages to the client that sent it") | |
| 1013 | 1021 | clean := pipeCmd.Bool("c", false, "Don't send status messages") | |
| 1022 | + | dispatcher := pipeCmd.String("d", "multicast", "Type of dispatcher (e.g. multicast, round_robin)") | |
| 1014 | 1023 | ||
| 1015 | 1024 | if !flagCheck(pipeCmd, topic, cmd.args) { | |
| 1016 | 1025 | return fmt.Errorf("invalid cmd args") |
| ... | ... | @@ -1028,6 +1037,7 @@ func (handler *CliHandler) pipe(cmd *CliCmd, topic string, clientID string) erro | |
| 1028 | 1037 | "topic", topic, | |
| 1029 | 1038 | "access", *access, | |
| 1030 | 1039 | "clean", *clean, | |
| 1040 | + | "dispatcher", *dispatcher, | |
| 1031 | 1041 | ) | |
| 1032 | 1042 | ||
| 1033 | 1043 | var accessList []string |
| ... | ... | @@ -1101,6 +1111,14 @@ func (handler *CliHandler) pipe(cmd *CliCmd, topic string, clientID string) erro | |
| 1101 | 1111 | ||
| 1102 | 1112 | throttledRW := newThrottledMonitorRW(cmd.sesh, handler, cmd, name) | |
| 1103 | 1113 | ||
| 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 | + | ||
| 1104 | 1122 | readErr, writeErr := handler.PubSub.Pipe( | |
| 1105 | 1123 | cmd.pipeCtx, | |
| 1106 | 1124 | clientID, |
+14,
-0
| ... | ... | @@ -5,6 +5,7 @@ import ( | |
| 5 | 5 | "io" | |
| 6 | 6 | "iter" | |
| 7 | 7 | "log/slog" | |
| 8 | + | "reflect" | |
| 8 | 9 | "sync" | |
| 9 | 10 | "time" | |
| 10 | 11 |
| ... | ... | @@ -21,6 +22,7 @@ type Broker interface { | |
| 21 | 22 | GetChannels() iter.Seq2[string, *Channel] | |
| 22 | 23 | GetClients() iter.Seq2[string, *Client] | |
| 23 | 24 | Connect(*Client, []*Channel) (error, error) | |
| 25 | + | SetDispatcher(dispatcher MessageDispatcher, channels []*Channel) error | |
| 24 | 26 | } | |
| 25 | 27 | ||
| 26 | 28 | type BaseBroker struct { |
| ... | ... | @@ -104,6 +106,7 @@ func (b *BaseBroker) Connect(client *Client, channels []*Channel) (error, error) | |
| 104 | 106 | for { | |
| 105 | 107 | data := make([]byte, 32*1024) | |
| 106 | 108 | n, err := client.ReadWriter.Read(data) | |
| 109 | + | ||
| 107 | 110 | data = data[:n] | |
| 108 | 111 | ||
| 109 | 112 | channelMessage := ChannelMessage{ |
| ... | ... | @@ -200,4 +203,15 @@ func (b *BaseBroker) ensureChannel(channel *Channel) *Channel { | |
| 200 | 203 | return dataChannel | |
| 201 | 204 | } | |
| 202 | 205 | ||
| 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 | + | ||
| 203 | 217 | var _ Broker = (*BaseBroker)(nil) |
+37,
-18
| ... | ... | @@ -57,12 +57,28 @@ type Channel struct { | |
| 57 | 57 | Clients *syncmap.Map[string, *Client] | |
| 58 | 58 | handleOnce sync.Once | |
| 59 | 59 | cleanupOnce sync.Once | |
| 60 | + | mu sync.Mutex | |
| 61 | + | Dispatcher MessageDispatcher | |
| 60 | 62 | } | |
| 61 | 63 | ||
| 62 | 64 | func (c *Channel) GetClients() iter.Seq2[string, *Client] { | |
| 63 | 65 | return c.Clients.Range | |
| 64 | 66 | } | |
| 65 | 67 | ||
| 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 | + | ||
| 66 | 82 | func (c *Channel) Cleanup() { | |
| 67 | 83 | c.cleanupOnce.Do(func() { | |
| 68 | 84 | close(c.Done) |
| ... | ... | @@ -83,30 +99,33 @@ func (c *Channel) Handle() { | |
| 83 | 99 | case <-c.Done: | |
| 84 | 100 | return | |
| 85 | 101 | 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) | |
| 87 | 114 | for _, client := range c.GetClients() { | |
| 115 | + | // Skip input-only clients and senders (unless replay is enabled) | |
| 88 | 116 | if client.Direction == ChannelDirectionInput || (client.ID == data.ClientID && !client.Replay) { | |
| 89 | 117 | continue | |
| 90 | 118 | } | |
| 119 | + | subscribers = append(subscribers, client) | |
| 120 | + | } | |
| 91 | 121 | ||
| 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 | + | } | |
| 108 | 128 | } | |
| 109 | - | wg.Wait() | |
| 110 | 129 | } | |
| 111 | 130 | } | |
| 112 | 131 | }() |
+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] { | |
| 58 | 58 | return p.getClients(ChannelDirectionOutput) | |
| 59 | 59 | } | |
| 60 | 60 | ||
| 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) { | |
| 62 | 62 | client := NewClient(ID, rw, direction, blockWrite, replay, keepAlive) | |
| 63 | 63 | ||
| 64 | + | // Set dispatcher on all channels (only if not already set) | |
| 65 | + | for _, ch := range channels { | |
| 66 | + | ch.SetDispatcher(dispatcher) | |
| 67 | + | } | |
| 68 | + | ||
| 64 | 69 | go func() { | |
| 65 | 70 | <-ctx.Done() | |
| 66 | 71 | client.Cleanup() |
| ... | ... | @@ -70,15 +75,15 @@ func (p *Multicast) connect(ctx context.Context, ID string, rw io.ReadWriter, ch | |
| 70 | 75 | } | |
| 71 | 76 | ||
| 72 | 77 | 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{}) | |
| 74 | 79 | } | |
| 75 | 80 | ||
| 76 | 81 | 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{})) | |
| 78 | 83 | } | |
| 79 | 84 | ||
| 80 | 85 | 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{})) | |
| 82 | 87 | } | |
| 83 | 88 | ||
| 84 | 89 | var _ PubSub = (*Multicast)(nil) |
+5,
-5
| ... | ... | @@ -66,7 +66,7 @@ func TestMulticastSubBlock(t *testing.T) { | |
| 66 | 66 | t.Fatalf("\norderActual:(%s)\norderExpected:(%s)", orderActual, orderExpected) | |
| 67 | 67 | } | |
| 68 | 68 | if actual.String() != expected { | |
| 69 | - | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual, expected) | |
| 69 | + | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual.String(), expected) | |
| 70 | 70 | } | |
| 71 | 71 | } | |
| 72 | 72 |
| ... | ... | @@ -96,8 +96,8 @@ func TestMulticastPubBlock(t *testing.T) { | |
| 96 | 96 | ||
| 97 | 97 | go func() { | |
| 98 | 98 | orderActual += "sub-" | |
| 99 | - | wg.Done() | |
| 100 | 99 | fmt.Println(cast.Sub(context.TODO(), "2", actual, []*Channel{channel}, false)) | |
| 100 | + | wg.Done() | |
| 101 | 101 | }() | |
| 102 | 102 | ||
| 103 | 103 | wg.Wait() |
| ... | ... | @@ -106,7 +106,7 @@ func TestMulticastPubBlock(t *testing.T) { | |
| 106 | 106 | t.Fatalf("\norderActual:(%s)\norderExpected:(%s)", orderActual, orderExpected) | |
| 107 | 107 | } | |
| 108 | 108 | if actual.String() != expected { | |
| 109 | - | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual, expected) | |
| 109 | + | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual.String(), expected) | |
| 110 | 110 | } | |
| 111 | 111 | } | |
| 112 | 112 |
| ... | ... | @@ -156,9 +156,9 @@ func TestMulticastMultSubs(t *testing.T) { | |
| 156 | 156 | t.Fatalf("\norderActual:(%s)\norderExpected:(%s)", orderActual, orderExpected) | |
| 157 | 157 | } | |
| 158 | 158 | if actual.String() != expected { | |
| 159 | - | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual, expected) | |
| 159 | + | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actual.String(), expected) | |
| 160 | 160 | } | |
| 161 | 161 | if actualOther.String() != expected { | |
| 162 | - | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actualOther, expected) | |
| 162 | + | t.Fatalf("\nactual:(%s)\nexpected:(%s)", actualOther.String(), expected) | |
| 163 | 163 | } | |
| 164 | 164 | } |
+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 | + | } |