135 lines
4.0 KiB
Go
135 lines
4.0 KiB
Go
package stream
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"sync/atomic"
|
|
)
|
|
|
|
const (
|
|
// subscriptionStateOpen is the default state of a subscription. An open
|
|
// subscription may receive new events.
|
|
subscriptionStateOpen uint32 = 0
|
|
|
|
// subscriptionStateClosed indicates that the subscription was closed, possibly
|
|
// as a result of a change to an ACL token, and will not receive new events.
|
|
// The subscriber must issue a new Subscribe request.
|
|
subscriptionStateClosed uint32 = 1
|
|
)
|
|
|
|
var (
|
|
// ErrSubscriptionReload is a error signalling reload event should be sent to
|
|
// the client and the server should close.
|
|
ErrSubscriptionReload = errors.New("subscription closed by server, client should retry")
|
|
)
|
|
|
|
// Subscription holds state about a single Subscribe call. Subscribe clients
|
|
// access their next event by calling Next(). This may initially include the
|
|
// snapshot events to catch them up if they are new or behind.
|
|
type Subscription struct {
|
|
// state is accessed atomically 0 means open, 1 means closed with reload
|
|
state uint32
|
|
|
|
// req is the requests that we are responding to
|
|
req *SubscribeRequest
|
|
|
|
// currentItem stores the current snapshot or topic buffer item we are on. It
|
|
// is mutated by calls to Next.
|
|
currentItem *BufferItem
|
|
|
|
// ctx is the Subscription context that wraps the context of the streaming RPC
|
|
// handler call.
|
|
ctx context.Context
|
|
|
|
// cancelFn stores the context cancel function that will wake up the
|
|
// in-progress Next call on a server-initiated state change e.g. Reload.
|
|
cancelFn func()
|
|
|
|
// Unsubscribe is a function set by EventPublisher that is called to
|
|
// free resources when the subscription is no longer needed.
|
|
Unsubscribe func()
|
|
}
|
|
|
|
// SubscribeRequest identifies the types of events the subscriber would like to
|
|
// receiver. Topic and Token are required.
|
|
type SubscribeRequest struct {
|
|
Topic Topic
|
|
Key string
|
|
Token string
|
|
Index uint64
|
|
}
|
|
|
|
// NewSubscription return a new subscription. The caller is responsible for
|
|
// calling Unsubscribe when it is done with the subscription, to free resources.
|
|
func NewSubscription(ctx context.Context, req *SubscribeRequest, item *BufferItem) *Subscription {
|
|
subCtx, cancel := context.WithCancel(ctx)
|
|
return &Subscription{
|
|
ctx: subCtx,
|
|
cancelFn: cancel,
|
|
req: req,
|
|
currentItem: item,
|
|
}
|
|
}
|
|
|
|
// Next returns the next set of events to deliver. It must only be called from a
|
|
// single goroutine concurrently as it mutates the Subscription.
|
|
func (s *Subscription) Next() ([]Event, error) {
|
|
if atomic.LoadUint32(&s.state) == subscriptionStateClosed {
|
|
return nil, ErrSubscriptionReload
|
|
}
|
|
|
|
for {
|
|
next, err := s.currentItem.Next(s.ctx)
|
|
if err != nil {
|
|
// Check we didn't return because of a state change cancelling the context
|
|
if atomic.LoadUint32(&s.state) == subscriptionStateClosed {
|
|
return nil, ErrSubscriptionReload
|
|
}
|
|
return nil, err
|
|
}
|
|
// Advance our cursor for next loop or next call
|
|
s.currentItem = next
|
|
|
|
// Assume happy path where all events (or none) are relevant.
|
|
allMatch := true
|
|
|
|
// If there is a specific key, see if we need to filter any events
|
|
if s.req.Key != "" {
|
|
for _, e := range next.Events {
|
|
if s.req.Key != e.Key {
|
|
allMatch = false
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
// Only if we need to filter events should we bother allocating a new slice
|
|
// as this is a hot loop.
|
|
events := next.Events
|
|
if !allMatch {
|
|
events = make([]Event, 0, len(next.Events))
|
|
for _, e := range next.Events {
|
|
// Only return it if the key matches.
|
|
if s.req.Key == "" || s.req.Key == e.Key {
|
|
events = append(events, e)
|
|
}
|
|
}
|
|
}
|
|
|
|
if len(events) > 0 {
|
|
return events, nil
|
|
}
|
|
// Keep looping until we find some events we are interested in.
|
|
}
|
|
}
|
|
|
|
// Close the subscription. Subscribers will receive an error when they call Next,
|
|
// and will need to perform a new Subscribe request.
|
|
// It is safe to call from any goroutine.
|
|
func (s *Subscription) Close() {
|
|
swapped := atomic.CompareAndSwapUint32(&s.state, subscriptionStateOpen, subscriptionStateClosed)
|
|
if swapped {
|
|
s.cancelFn()
|
|
}
|
|
}
|