diff --git a/rpc/stream/cond.go b/rpc/stream/cond.go index 8f629907ee..29f902ce28 100644 --- a/rpc/stream/cond.go +++ b/rpc/stream/cond.go @@ -15,11 +15,19 @@ func NewCond() *Cond { return &Cond{ch: make(chan struct{})} } +func (c *Cond) Channel() <-chan struct{} { + c.mu.Lock() + defer c.mu.Unlock() + + return c.ch +} + // Wait returns true if the condition is signaled, false if the context is canceled func (c *Cond) Wait(ctx context.Context) bool { - c.mu.Lock() - ch := c.ch - c.mu.Unlock() + return c.WaitChannel(ctx, c.Channel()) +} + +func (c *Cond) WaitChannel(ctx context.Context, ch <-chan struct{}) bool { select { case <-ch: diff --git a/rpc/stream/stream.go b/rpc/stream/stream.go index 794bcb2d5a..5bd36ca7bf 100644 --- a/rpc/stream/stream.go +++ b/rpc/stream/stream.go @@ -71,8 +71,13 @@ func (s *Stream[V]) Add(vs ...V) int { func (s *Stream[V]) Subscribe(ctx context.Context, callback func([]V, int) error) error { var ( items []V - offset = -1 + offset int ) + + s.mutex.RLock() + offset = s.lastID() + s.mutex.RUnlock() + for { items, offset = s.ReadBlocking(ctx, offset) if len(items) == 0 { @@ -130,8 +135,9 @@ func (s *Stream[V]) ReadBlocking(ctx context.Context, offset int) ([]V, int) { return items, offset } + ch := s.cond.Channel() s.mutex.RUnlock() - r := s.cond.Wait(ctx) + r := s.cond.WaitChannel(ctx, ch) s.mutex.RLock() if !r {