Commit cc428a5

Antonio Mika  ·  2026-10-10 22:14:55 -0400 EDT
parent 85628be 
fix(pubsub): deliver a publisher's last message before closing subscribers
3 files changed,  +70, -1
M pkg/pubsub/broker.go
M pkg/pubsub/client.go
M pkg/pubsub/regression_test.go
M pkg/pubsub/broker.go
+4, -1
......@@ -129,7 +129,7 @@ func (b *BaseBroker) Connect(client *Client, channels []*Channel) (error, error)
129129 }
130130 }
131131 if otherChannels == 0 {
132- cl.Cleanup()
132+ cl.Finish()
133133 }
134134 }
135135 }
......@@ -236,6 +236,9 @@ func (b *BaseBroker) Connect(client *Client, channels []*Channel) (error, error)
236236 }
237237 case <-client.Done:
238238 break mainLoop
239+ case <-client.finish:
240+ client.Cleanup()
241+ break mainLoop
239242 }
240243 }
241244 }()
M pkg/pubsub/client.go
+11, -0
......@@ -15,6 +15,7 @@ func NewClient(ID string, rw io.ReadWriter, direction ChannelDirection, blockWri
1515 Direction: direction,
1616 Channels: syncmap.New[string, *Channel](),
1717 Done: make(chan struct{}),
18+ finish: make(chan struct{}),
1819 Data: make(chan ChannelMessage),
1920 Replay: replay,
2021 BlockWrite: blockWrite,
......@@ -39,6 +40,8 @@ type Client struct {
3940 KeepAlive bool
4041 once sync.Once
4142 onceData sync.Once
43+ finish chan struct{}
44+ finishOnce sync.Once
4245 }
4346
4447 func (c *Client) GetChannels() iter.Seq2[string, *Channel] {
......@@ -53,3 +56,11 @@ func (c *Client) Cleanup() {
5356 }
5457 })
5558 }
59+
60+// Finish ends a subscriber after it writes the message it was last sent, so
61+// a publisher that leaves doesn't cut off its own last message.
62+func (c *Client) Finish() {
63+ c.finishOnce.Do(func() {
64+ close(c.finish)
65+ })
66+}
M pkg/pubsub/regression_test.go
+55, -0
......@@ -4,6 +4,7 @@ import (
44 "bytes"
55 "context"
66 "fmt"
7+ "io"
78 "log/slog"
89 "sync"
910 "testing"
......@@ -221,3 +222,57 @@ func TestDispatcherSingleSubscriber(t *testing.T) {
221222 }
222223 })
223224 }
225+
226+// slowSession is a subscriber whose writes take a while and fail once it is
227+// closed, like an ssh session.
228+type slowSession struct {
229+ Buffer
230+ mu sync.Mutex
231+ closed bool
232+}
233+
234+func (s *slowSession) Write(p []byte) (int, error) {
235+ time.Sleep(20 * time.Millisecond)
236+ s.mu.Lock()
237+ defer s.mu.Unlock()
238+ if s.closed {
239+ return 0, io.ErrClosedPipe
240+ }
241+ return s.Buffer.Write(p)
242+}
243+
244+func (s *slowSession) Close() error {
245+ s.mu.Lock()
246+ defer s.mu.Unlock()
247+ s.closed = true
248+ return nil
249+}
250+
251+func TestSubReceivesLastMessageWhenPublisherLeaves(t *testing.T) {
252+ cast := NewMulticast(slog.Default())
253+ ctx, cancel := context.WithCancel(context.Background())
254+ defer cancel()
255+
256+ sub := &slowSession{}
257+ subDone := make(chan struct{})
258+ go func() {
259+ _ = cast.Sub(ctx, "sub", sub, []*Channel{NewChannel("topic")}, false)
260+ close(subDone)
261+ }()
262+ for joined := 0; joined == 0; time.Sleep(time.Millisecond) {
263+ for range cast.GetSubs() {
264+ joined++
265+ }
266+ }
267+
268+ _ = cast.Pub(ctx, "pub", &Buffer{b: *bytes.NewBufferString("last")}, []*Channel{NewChannel("topic")}, false)
269+
270+ select {
271+ case <-subDone:
272+ case <-time.After(5 * time.Second):
273+ t.Fatal("subscriber is still connected after the publisher left")
274+ }
275+ if got := sub.String(); got != "last" {
276+ t.Errorf("subscriber got %q, want %q", got, "last")
277+ }
278+}