hanzoai/psrpc
2
0
Fork 0
mirror of https://github.com/hanzoai/psrpc.git synced 2026-09-30 23:47:22 +00:00
Code Issues Packages Projects Releases Wiki Activity

unblock subscriptions on close (#94)

Browse Source
This commit is contained in:
Paul Wells
2024-04-25 21:50:48 -07:00
committed by GitHub
parent 59837372c0
commit cb3264b945
3 changed files with 91 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
internal/bus/bus_local.go
+33 -19
View File
@@ -54,15 +54,15 @@ func (l *localMessageBus) Publish(_ context.Context, channel Channel, msg proto.
return nil
}
func (l *localMessageBus) Subscribe(_ context.Context, channel Channel, size int) (Reader, error) {
return l.subscribe(l.subs, channel.Legacy, size, false)
func (l *localMessageBus) Subscribe(ctx context.Context, channel Channel, size int) (Reader, error) {
return l.subscribe(ctx, l.subs, channel.Legacy, size, false)
}
func (l *localMessageBus) SubscribeQueue(_ context.Context, channel Channel, size int) (Reader, error) {
return l.subscribe(l.queues, channel.Legacy, size, true)
func (l *localMessageBus) SubscribeQueue(ctx context.Context, channel Channel, size int) (Reader, error) {
return l.subscribe(ctx, l.queues, channel.Legacy, size, true)
}
func (l *localMessageBus) subscribe(subLists map[string]*localSubList, channel string, size int, queue bool) (Reader, error) {
func (l *localMessageBus) subscribe(ctx context.Context, subLists map[string]*localSubList, channel string, size int, queue bool) (Reader, error) {
l.Lock()
defer l.Unlock()
@@ -74,7 +74,6 @@ func (l *localMessageBus) subscribe(subLists map[string]*localSubList, channel s
l.Lock()
subList.Lock()
close(subList.subs[index])
subList.subs[index] = nil
subList.subCount--
if subList.subCount == 0 {
@@ -87,20 +86,25 @@ func (l *localMessageBus) subscribe(subLists map[string]*localSubList, channel s
subLists[channel] = subList
}
return subList.create(size), nil
return subList.create(ctx, size), nil
}
type localSubList struct {
sync.RWMutex // locking while holding localMessageBus lock is allowed
subs []chan []byte
subs []*localSubscription
subCount int
queue bool
next int
onUnsubscribe func(int)
}
func (l *localSubList) create(size int) *localSubscription {
msgChan := make(chan []byte, size)
func (l *localSubList) create(ctx context.Context, size int) *localSubscription {
ctx, cancel := context.WithCancel(ctx)
sub := &localSubscription{
ctx: ctx,
cancel: cancel,
msgChan: make(chan []byte, size),
}
l.Lock()
defer l.Unlock()
@@ -112,22 +116,21 @@ func (l *localSubList) create(size int) *localSubscription {
if s == nil {
added = true
index = i
l.subs[i] = msgChan
l.subs[i] = sub
break
}
}
if !added {
index = len(l.subs)
l.subs = append(l.subs, msgChan)
l.subs = append(l.subs, sub)
}
return &localSubscription{
msgChan: msgChan,
onClose: func() {
l.onUnsubscribe(index)
},
sub.onClose = func() {
l.onUnsubscribe(index)
}
return sub
}
func (l *localSubList) dispatch(b []byte) {
@@ -143,7 +146,7 @@ func (l *localSubList) dispatch(b []byte) {
s := l.subs[l.next]
l.next++
if s != nil {
s <- b
s.write(b)
return
}
}
@@ -154,17 +157,26 @@ func (l *localSubList) dispatch(b []byte) {
// send to all
for _, s := range l.subs {
if s != nil {
s <- b
s.write(b)
}
}
}
}
type localSubscription struct {
ctx context.Context
cancel context.CancelFunc
msgChan chan []byte
onClose func()
}
func (l *localSubscription) write(b []byte) {
select {
case l.msgChan <- b:
case <-l.ctx.Done():
}
}
func (l *localSubscription) read() ([]byte, bool) {
msg, ok := <-l.msgChan
if !ok {
@@ -174,6 +186,8 @@ func (l *localSubscription) read() ([]byte, bool) {
}
func (l *localSubscription) Close() error {
l.cancel()
l.onClose()
close(l.msgChan)
return nil
}
internal/bus/bus_nats.go
+38 -17
View File
@@ -46,39 +46,41 @@ func (n *natsMessageBus) Publish(_ context.Context, channel Channel, msg proto.M
return n.nc.Publish(channel.Server, b)
}
func (n *natsMessageBus) Subscribe(_ context.Context, channel Channel, size int) (Reader, error) {
func (n *natsMessageBus) Subscribe(ctx context.Context, channel Channel, size int) (Reader, error) {
if channel.Local == "" {
return n.subscribe(channel.Server, size, false)
return n.subscribe(ctx, channel.Server, size, false)
} else {
return n.subscribeRouter(channel, size, false)
return n.subscribeRouter(ctx, channel, size, false)
}
}
func (n *natsMessageBus) SubscribeQueue(_ context.Context, channel Channel, size int) (Reader, error) {
func (n *natsMessageBus) SubscribeQueue(ctx context.Context, channel Channel, size int) (Reader, error) {
if channel.Local == "" {
return n.subscribe(channel.Server, size, true)
return n.subscribe(ctx, channel.Server, size, true)
} else {
return n.subscribeRouter(channel, size, true)
return n.subscribeRouter(ctx, channel, size, true)
}
}
func (n *natsMessageBus) subscribe(channel string, size int, queue bool) (*natsSubscription, error) {
msgChan := make(chan *nats.Msg, size)
var sub *nats.Subscription
func (n *natsMessageBus) subscribe(ctx context.Context, channel string, size int, queue bool) (*natsSubscription, error) {
ctx, cancel := context.WithCancel(ctx)
sub := &natsSubscription{
ctx: ctx,
cancel: cancel,
msgChan: make(chan *nats.Msg, size),
}
var err error
if queue {
sub, err = n.nc.ChanQueueSubscribe(channel, "bus", msgChan)
sub.sub, err = n.nc.QueueSubscribe(channel, "bus", sub.write)
} else {
sub, err = n.nc.ChanSubscribe(channel, msgChan)
sub.sub, err = n.nc.Subscribe(channel, sub.write)
}
if err != nil {
return nil, err
}
return &natsSubscription{
sub: sub,
msgChan: msgChan,
}, nil
return sub, nil
}
func (n *natsMessageBus) unsubscribeRouter(r *natsRouter, channel string, s *natsRouterSubscription) {
@@ -89,8 +91,11 @@ func (n *natsMessageBus) unsubscribeRouter(r *natsRouter, channel string, s *nat
}
}
func (n *natsMessageBus) subscribeRouter(channel Channel, size int, queue bool) (*natsRouterSubscription, error) {
func (n *natsMessageBus) subscribeRouter(ctx context.Context, channel Channel, size int, queue bool) (*natsRouterSubscription, error) {
ctx, cancel := context.WithCancel(ctx)
sub := &natsRouterSubscription{
ctx: ctx,
cancel: cancel,
msgChan: make(chan *nats.Msg, size),
channel: channel.Local,
}
@@ -135,10 +140,19 @@ func (n *natsMessageBus) subscribeRouter(channel Channel, size int, queue bool)
}
type natsSubscription struct {
ctx context.Context
cancel context.CancelFunc
sub *nats.Subscription
msgChan chan *nats.Msg
}
func (n *natsSubscription) write(msg *nats.Msg) {
select {
case n.msgChan <- msg:
case <-n.ctx.Done():
}
}
func (n *natsSubscription) read() ([]byte, bool) {
msg, ok := <-n.msgChan
if !ok {
@@ -148,6 +162,7 @@ func (n *natsSubscription) read() ([]byte, bool) {
}
func (n *natsSubscription) Close() error {
n.cancel()
err := n.sub.Unsubscribe()
close(n.msgChan)
return err
@@ -205,13 +220,18 @@ func (n *natsRouter) write(m *nats.Msg) {
}
type natsRouterSubscription struct {
ctx context.Context
cancel context.CancelFunc
msgChan chan *nats.Msg
router *natsRouter
channel string
}
func (n *natsRouterSubscription) write(m *nats.Msg) {
n.msgChan <- m
select {
case n.msgChan <- m:
case <-n.ctx.Done():
}
}
func (n *natsRouterSubscription) read() ([]byte, bool) {
@@ -223,6 +243,7 @@ func (n *natsRouterSubscription) read() ([]byte, bool) {
}
func (n *natsRouterSubscription) Close() error {
n.cancel()
n.router.bus.unsubscribeRouter(n.router, n.channel, n)
close(n.msgChan)
return nil
internal/bus/bus_redis.go
+20 -10
View File
@@ -103,9 +103,11 @@ func (r *redisMessageBus) SubscribeQueue(ctx context.Context, channel Channel, s
}
func (r *redisMessageBus) subscribe(ctx context.Context, channel string, size int, subLists map[string]*redisSubList, queue bool) (Reader, error) {
ctx, cancel := context.WithCancel(ctx)
sub := &redisSubscription{
bus: r,
ctx: ctx,
cancel: cancel,
channel: channel,
msgChan: make(chan *redis.Message, size),
queue: queue,
@@ -120,12 +122,12 @@ func (r *redisMessageBus) subscribe(ctx context.Context, channel string, size in
subLists[channel] = subList
r.reconcileSubscriptions(channel)
}
subList.subs = append(subList.subs, sub.msgChan)
subList.subs = append(subList.subs, sub)
return sub, nil
}
func (r *redisMessageBus) unsubscribe(channel string, queue bool, msgChan chan *redis.Message) {
func (r *redisMessageBus) unsubscribe(channel string, queue bool, sub *redisSubscription) {
r.mu.Lock()
defer r.mu.Unlock()
@@ -140,13 +142,12 @@ func (r *redisMessageBus) unsubscribe(channel string, queue bool, msgChan chan *
if !ok {
return
}
i := slices.Index(subList.subs, msgChan)
i := slices.Index(subList.subs, sub)
if i == -1 {
return
}
subList.subs = slices.Delete(subList.subs, i, i+1)
close(msgChan)
if len(subList.subs) == 0 {
delete(subLists, channel)
@@ -320,7 +321,7 @@ func (r *redisReconcileSubscriptionsOp) run() error {
}
type redisSubList struct {
subs []chan *redis.Message
subs []*redisSubscription
next int
}
@@ -328,24 +329,32 @@ func (r *redisSubList) dispatchQueue(msg *redis.Message) {
if r.next >= len(r.subs) {
r.next = 0
}
r.subs[r.next] <- msg
r.subs[r.next].write(msg)
r.next++
}
func (r *redisSubList) dispatch(msg *redis.Message) {
for _, ch := range r.subs {
ch <- msg
for _, sub := range r.subs {
sub.write(msg)
}
}
type redisSubscription struct {
bus *redisMessageBus
ctx context.Context
cancel context.CancelFunc
channel string
msgChan chan *redis.Message
queue bool
}
func (r *redisSubscription) write(msg *redis.Message) {
select {
case r.msgChan <- msg:
case <-r.ctx.Done():
}
}
func (r *redisSubscription) read() ([]byte, bool) {
for {
var msg *redis.Message
@@ -356,7 +365,6 @@ func (r *redisSubscription) read() ([]byte, bool) {
return nil, false
}
case <-r.ctx.Done():
r.Close()
return nil, false
}
@@ -374,6 +382,8 @@ func (r *redisSubscription) read() ([]byte, bool) {
}
func (r *redisSubscription) Close() error {
r.bus.unsubscribe(r.channel, r.queue, r.msgChan)
r.cancel()
r.bus.unsubscribe(r.channel, r.queue, r)
close(r.msgChan)
return nil
}