OP#792 initial commit after fork
This commit is contained in:
66
internal/pubsub/mask.go
Normal file
66
internal/pubsub/mask.go
Normal file
@@ -0,0 +1,66 @@
|
||||
package pubsub
|
||||
|
||||
import (
|
||||
"math"
|
||||
"math/bits"
|
||||
)
|
||||
|
||||
// Mask allows filtering by a bitset mask.
|
||||
type Mask uint64
|
||||
|
||||
const (
|
||||
// MaskAll is the mask for all entries.
|
||||
MaskAll Mask = math.MaxUint64
|
||||
)
|
||||
|
||||
// MaskFromMaskable extracts mask from an interface.
|
||||
func MaskFromMaskable(m Maskable) Mask {
|
||||
return Mask(m.Mask())
|
||||
}
|
||||
|
||||
// Contains returns whether *all* flags in other is present in t.
|
||||
func (t Mask) Contains(other Mask) bool {
|
||||
return t&other == other
|
||||
}
|
||||
|
||||
// Overlaps returns whether *any* flags in t overlaps with other.
|
||||
func (t Mask) Overlaps(other Mask) bool {
|
||||
return t&other != 0
|
||||
}
|
||||
|
||||
// SingleType returns whether t has a single type set.
|
||||
func (t Mask) SingleType() bool {
|
||||
return bits.OnesCount64(uint64(t)) == 1
|
||||
}
|
||||
|
||||
// FromUint64 will set a mask to the uint64 value.
|
||||
func (t *Mask) FromUint64(m uint64) {
|
||||
*t = Mask(m)
|
||||
}
|
||||
|
||||
// Merge will merge other into t.
|
||||
func (t *Mask) Merge(other Mask) {
|
||||
*t |= other
|
||||
}
|
||||
|
||||
// MergeMaskable will merge other into t.
|
||||
func (t *Mask) MergeMaskable(other Maskable) {
|
||||
*t |= Mask(other.Mask())
|
||||
}
|
||||
|
||||
// SetIf will add other if b is true.
|
||||
func (t *Mask) SetIf(b bool, other Mask) {
|
||||
if b {
|
||||
*t |= other
|
||||
}
|
||||
}
|
||||
|
||||
// Mask returns the mask as a uint64.
|
||||
func (t Mask) Mask() uint64 {
|
||||
return uint64(t)
|
||||
}
|
||||
|
||||
// Maskable implementations must return their mask as a 64 bit uint.
|
||||
type Maskable interface {
|
||||
Mask() uint64
|
||||
}
|
||||
196
internal/pubsub/pubsub.go
Normal file
196
internal/pubsub/pubsub.go
Normal file
@@ -0,0 +1,196 @@
|
||||
// Copyright (c) 2015-2024 MinIO, Inc.
|
||||
//
|
||||
// This file is part of MinIO Object Storage stack
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
package pubsub
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
)
|
||||
|
||||
// GetByteBuffer returns a byte buffer from the pool.
|
||||
var GetByteBuffer = func() []byte {
|
||||
return make([]byte, 0, 4096)
|
||||
}
|
||||
|
||||
// Sub - subscriber entity.
|
||||
type Sub[T Maskable] struct {
|
||||
ch chan T
|
||||
types Mask
|
||||
filter func(entry T) bool
|
||||
}
|
||||
|
||||
// PubSub holds publishers and subscribers
|
||||
type PubSub[T Maskable, M Maskable] struct {
|
||||
// atomics, keep at top:
|
||||
types uint64
|
||||
numSubscribers int32
|
||||
maxSubscribers int32
|
||||
|
||||
// not atomics:
|
||||
subs []*Sub[T]
|
||||
sync.RWMutex
|
||||
}
|
||||
|
||||
// Publish message to the subscribers.
|
||||
// Note that publish is always non-blocking send so that we don't block on slow receivers.
|
||||
// Hence receivers should use buffered channel so as not to miss the published events.
|
||||
func (ps *PubSub[T, M]) Publish(item T) {
|
||||
ps.RLock()
|
||||
defer ps.RUnlock()
|
||||
for _, sub := range ps.subs {
|
||||
if sub.types.Contains(Mask(item.Mask())) && (sub.filter == nil || sub.filter(item)) {
|
||||
select {
|
||||
case sub.ch <- item:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Subscribe - Adds a subscriber to pubsub system
|
||||
func (ps *PubSub[T, M]) Subscribe(mask M, subCh chan T, doneCh <-chan struct{}, filter func(entry T) bool) error {
|
||||
totalSubs := atomic.AddInt32(&ps.numSubscribers, 1)
|
||||
if ps.maxSubscribers > 0 && totalSubs > ps.maxSubscribers {
|
||||
atomic.AddInt32(&ps.numSubscribers, -1)
|
||||
return fmt.Errorf("the limit of `%d` subscribers is reached", ps.maxSubscribers)
|
||||
}
|
||||
ps.Lock()
|
||||
defer ps.Unlock()
|
||||
|
||||
sub := &Sub[T]{ch: subCh, types: Mask(mask.Mask()), filter: filter}
|
||||
ps.subs = append(ps.subs, sub)
|
||||
|
||||
// We hold a lock, so we are safe to update
|
||||
combined := Mask(atomic.LoadUint64(&ps.types))
|
||||
combined.Merge(Mask(mask.Mask()))
|
||||
atomic.StoreUint64(&ps.types, uint64(combined))
|
||||
|
||||
go func() {
|
||||
<-doneCh
|
||||
|
||||
ps.Lock()
|
||||
defer ps.Unlock()
|
||||
var remainTypes Mask
|
||||
for i, s := range ps.subs {
|
||||
if s == sub {
|
||||
ps.subs = append(ps.subs[:i], ps.subs[i+1:]...)
|
||||
} else {
|
||||
remainTypes.Merge(s.types)
|
||||
}
|
||||
}
|
||||
atomic.StoreUint64(&ps.types, uint64(remainTypes))
|
||||
atomic.AddInt32(&ps.numSubscribers, -1)
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// SubscribeJSON - Adds a subscriber to pubsub system and returns results with JSON encoding.
|
||||
func (ps *PubSub[T, M]) SubscribeJSON(mask M, subCh chan<- []byte, doneCh <-chan struct{}, filter func(entry T) bool, wg *sync.WaitGroup) error {
|
||||
totalSubs := atomic.AddInt32(&ps.numSubscribers, 1)
|
||||
if ps.maxSubscribers > 0 && totalSubs > ps.maxSubscribers {
|
||||
atomic.AddInt32(&ps.numSubscribers, -1)
|
||||
return fmt.Errorf("the limit of `%d` subscribers is reached", ps.maxSubscribers)
|
||||
}
|
||||
ps.Lock()
|
||||
defer ps.Unlock()
|
||||
subChT := make(chan T, 10000)
|
||||
sub := &Sub[T]{ch: subChT, types: Mask(mask.Mask()), filter: filter}
|
||||
ps.subs = append(ps.subs, sub)
|
||||
|
||||
// We hold a lock, so we are safe to update
|
||||
combined := Mask(atomic.LoadUint64(&ps.types))
|
||||
combined.Merge(Mask(mask.Mask()))
|
||||
atomic.StoreUint64(&ps.types, uint64(combined))
|
||||
if wg != nil {
|
||||
wg.Add(1)
|
||||
}
|
||||
go func() {
|
||||
defer func() {
|
||||
if wg != nil {
|
||||
wg.Done()
|
||||
}
|
||||
// Clean up and de-register the subscriber
|
||||
ps.Lock()
|
||||
defer ps.Unlock()
|
||||
var remainTypes Mask
|
||||
for i, s := range ps.subs {
|
||||
if s == sub {
|
||||
ps.subs = append(ps.subs[:i], ps.subs[i+1:]...)
|
||||
} else {
|
||||
remainTypes.Merge(s.types)
|
||||
}
|
||||
}
|
||||
atomic.StoreUint64(&ps.types, uint64(remainTypes))
|
||||
atomic.AddInt32(&ps.numSubscribers, -1)
|
||||
}()
|
||||
|
||||
// Read from subChT and write to subCh
|
||||
var buf bytes.Buffer
|
||||
enc := json.NewEncoder(&buf)
|
||||
for {
|
||||
select {
|
||||
case <-doneCh:
|
||||
return
|
||||
case v, ok := <-subChT:
|
||||
if !ok {
|
||||
return
|
||||
}
|
||||
buf.Reset()
|
||||
err := enc.Encode(v)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
|
||||
select {
|
||||
case subCh <- append(GetByteBuffer()[:0], buf.Bytes()...):
|
||||
continue
|
||||
case <-doneCh:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// NumSubscribers returns the number of current subscribers,
|
||||
// The mask is checked against the active subscribed types,
|
||||
// and 0 will be returned if nobody is subscribed for the type(s).
|
||||
func (ps *PubSub[T, M]) NumSubscribers(mask M) int32 {
|
||||
types := Mask(atomic.LoadUint64(&ps.types))
|
||||
if !types.Overlaps(Mask(mask.Mask())) {
|
||||
return 0
|
||||
}
|
||||
return atomic.LoadInt32(&ps.numSubscribers)
|
||||
}
|
||||
|
||||
// Subscribers returns the number of current subscribers for all types.
|
||||
func (ps *PubSub[T, M]) Subscribers() int32 {
|
||||
return atomic.LoadInt32(&ps.numSubscribers)
|
||||
}
|
||||
|
||||
// New inits a PubSub system with a limit of maximum
|
||||
// subscribers unless zero is specified
|
||||
func New[T Maskable, M Maskable](maxSubscribers int32) *PubSub[T, M] {
|
||||
return &PubSub[T, M]{maxSubscribers: maxSubscribers}
|
||||
}
|
||||
199
internal/pubsub/pubsub_test.go
Normal file
199
internal/pubsub/pubsub_test.go
Normal file
@@ -0,0 +1,199 @@
|
||||
// Copyright (c) 2015-2021 MinIO, Inc.
|
||||
//
|
||||
// This file is part of MinIO Object Storage stack
|
||||
//
|
||||
// This program is free software: you can redistribute it and/or modify
|
||||
// it under the terms of the GNU Affero General Public License as published by
|
||||
// the Free Software Foundation, either version 3 of the License, or
|
||||
// (at your option) any later version.
|
||||
//
|
||||
// This program is distributed in the hope that it will be useful
|
||||
// but WITHOUT ANY WARRANTY; without even the implied warranty of
|
||||
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
||||
// GNU Affero General Public License for more details.
|
||||
//
|
||||
// You should have received a copy of the GNU Affero General Public License
|
||||
// along with this program. If not, see <http://www.gnu.org/licenses/>.
|
||||
|
||||
package pubsub
|
||||
|
||||
import (
|
||||
"testing"
|
||||
"time"
|
||||
)
|
||||
|
||||
func TestSubscribe(t *testing.T) {
|
||||
ps := New[Maskable, Mask](2)
|
||||
ch1 := make(chan Maskable, 1)
|
||||
ch2 := make(chan Maskable, 1)
|
||||
doneCh := make(chan struct{})
|
||||
defer close(doneCh)
|
||||
if err := ps.Subscribe(MaskAll, ch1, doneCh, nil); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if err := ps.Subscribe(MaskAll, ch2, doneCh, nil); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
ps.Lock()
|
||||
defer ps.Unlock()
|
||||
|
||||
if len(ps.subs) != 2 || ps.NumSubscribers(MaskAll) != 2 || ps.Subscribers() != 2 {
|
||||
t.Fatalf("expected 2 subscribers")
|
||||
}
|
||||
}
|
||||
|
||||
func TestNumSubscribersMask(t *testing.T) {
|
||||
ps := New[Maskable, Mask](2)
|
||||
ch1 := make(chan Maskable, 1)
|
||||
ch2 := make(chan Maskable, 1)
|
||||
doneCh := make(chan struct{})
|
||||
defer close(doneCh)
|
||||
if err := ps.Subscribe(Mask(1), ch1, doneCh, nil); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if err := ps.Subscribe(Mask(2), ch2, doneCh, nil); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
ps.Lock()
|
||||
defer ps.Unlock()
|
||||
|
||||
if len(ps.subs) != 2 {
|
||||
t.Fatalf("expected 2 subscribers")
|
||||
}
|
||||
if want, got := int32(2), ps.NumSubscribers(Mask(1)); got != want {
|
||||
t.Fatalf("want %d subscribers, got %d", want, got)
|
||||
}
|
||||
if want, got := int32(2), ps.NumSubscribers(Mask(2)); got != want {
|
||||
t.Fatalf("want %d subscribers, got %d", want, got)
|
||||
}
|
||||
if want, got := int32(2), ps.NumSubscribers(Mask(1|2)); got != want {
|
||||
t.Fatalf("want %d subscribers, got %d", want, got)
|
||||
}
|
||||
if want, got := int32(2), ps.NumSubscribers(MaskAll); got != want {
|
||||
t.Fatalf("want %d subscribers, got %d", want, got)
|
||||
}
|
||||
if want, got := int32(0), ps.NumSubscribers(Mask(4)); got != want {
|
||||
t.Fatalf("want %d subscribers, got %d", want, got)
|
||||
}
|
||||
}
|
||||
|
||||
func TestSubscribeExceedingLimit(t *testing.T) {
|
||||
ps := New[Maskable, Maskable](2)
|
||||
ch1 := make(chan Maskable, 1)
|
||||
ch2 := make(chan Maskable, 1)
|
||||
ch3 := make(chan Maskable, 1)
|
||||
doneCh := make(chan struct{})
|
||||
defer close(doneCh)
|
||||
if err := ps.Subscribe(MaskAll, ch1, doneCh, nil); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if err := ps.Subscribe(MaskAll, ch2, doneCh, nil); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if err := ps.Subscribe(MaskAll, ch3, doneCh, nil); err == nil {
|
||||
t.Fatalf("unexpected nil err")
|
||||
}
|
||||
}
|
||||
|
||||
func TestUnsubscribe(t *testing.T) {
|
||||
ps := New[Maskable, Maskable](2)
|
||||
ch1 := make(chan Maskable, 1)
|
||||
ch2 := make(chan Maskable, 1)
|
||||
doneCh1 := make(chan struct{})
|
||||
doneCh2 := make(chan struct{})
|
||||
if err := ps.Subscribe(MaskAll, ch1, doneCh1, nil); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if err := ps.Subscribe(MaskAll, ch2, doneCh2, nil); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
|
||||
close(doneCh1)
|
||||
// Allow for the above statement to take effect.
|
||||
time.Sleep(100 * time.Millisecond)
|
||||
ps.Lock()
|
||||
if len(ps.subs) != 1 {
|
||||
t.Fatal("expected 1 subscriber")
|
||||
}
|
||||
ps.Unlock()
|
||||
close(doneCh2)
|
||||
}
|
||||
|
||||
type maskString string
|
||||
|
||||
func (m maskString) Mask() uint64 {
|
||||
return 1
|
||||
}
|
||||
|
||||
func TestPubSub(t *testing.T) {
|
||||
ps := New[Maskable, Maskable](1)
|
||||
ch1 := make(chan Maskable, 1)
|
||||
doneCh1 := make(chan struct{})
|
||||
defer close(doneCh1)
|
||||
if err := ps.Subscribe(MaskAll, ch1, doneCh1, func(entry Maskable) bool { return true }); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
val := maskString("hello")
|
||||
ps.Publish(val)
|
||||
msg := <-ch1
|
||||
if msg != val {
|
||||
t.Fatalf("expected %s , found %s", val, msg)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultiPubSub(t *testing.T) {
|
||||
ps := New[Maskable, Maskable](2)
|
||||
ch1 := make(chan Maskable, 1)
|
||||
ch2 := make(chan Maskable, 1)
|
||||
doneCh := make(chan struct{})
|
||||
defer close(doneCh)
|
||||
if err := ps.Subscribe(MaskAll, ch1, doneCh, func(entry Maskable) bool { return true }); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
if err := ps.Subscribe(MaskAll, ch2, doneCh, func(entry Maskable) bool { return true }); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
val := maskString("hello")
|
||||
ps.Publish(val)
|
||||
|
||||
msg1 := <-ch1
|
||||
msg2 := <-ch2
|
||||
if msg1 != val && msg2 != val {
|
||||
t.Fatalf("expected both subscribers to have%s , found %s and %s", val, msg1, msg2)
|
||||
}
|
||||
}
|
||||
|
||||
func TestMultiPubSubMask(t *testing.T) {
|
||||
ps := New[Maskable, Maskable](3)
|
||||
ch1 := make(chan Maskable, 1)
|
||||
ch2 := make(chan Maskable, 1)
|
||||
ch3 := make(chan Maskable, 1)
|
||||
doneCh := make(chan struct{})
|
||||
defer close(doneCh)
|
||||
// Mask matches maskString, should get result
|
||||
if err := ps.Subscribe(Mask(1), ch1, doneCh, func(entry Maskable) bool { return true }); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
// Mask matches maskString, should get result
|
||||
if err := ps.Subscribe(Mask(1|2), ch2, doneCh, func(entry Maskable) bool { return true }); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
// Does NOT overlap maskString
|
||||
if err := ps.Subscribe(Mask(2), ch3, doneCh, func(entry Maskable) bool { return true }); err != nil {
|
||||
t.Fatalf("unexpected error: %v", err)
|
||||
}
|
||||
val := maskString("hello")
|
||||
ps.Publish(val)
|
||||
|
||||
msg1 := <-ch1
|
||||
msg2 := <-ch2
|
||||
if msg1 != val && msg2 != val {
|
||||
t.Fatalf("expected both subscribers to have%s , found %s and %s", val, msg1, msg2)
|
||||
}
|
||||
|
||||
select {
|
||||
case msg := <-ch3:
|
||||
t.Fatalf("unexpected msg, f got %s", msg)
|
||||
default:
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user