snapshot of https://github.com/centrifugal/centrifugo.git (master, 2026-09-16; blobs > 131072B omitted)

internal/natsbroker/broker.go

@@ -0,0 +1,439 @@
+// Package natsbroker defines custom Nats Broker for Centrifuge library.
+package natsbroker
+
+import (
+	"context"
+	"fmt"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/centrifugal/centrifugo/v6/internal/configtypes"
+	"github.com/centrifugal/centrifugo/v6/internal/tools"
+
+	"github.com/centrifugal/centrifuge"
+	"github.com/centrifugal/protocol"
+	"github.com/nats-io/nats.go"
+	"github.com/rs/zerolog/log"
+)
+
+type (
+	// channelID is unique channel identifier in Nats.
+	channelID string
+)
+
+type Config = configtypes.NatsBroker
+
+type subWrapper struct {
+	sub         *nats.Subscription
+	origChannel string
+}
+
+// NatsBroker is a broker on top of Nats messaging system.
+type NatsBroker struct {
+	node   *centrifuge.Node
+	config Config
+
+	initConnOnce        sync.Once
+	nc                  *nats.Conn
+	subsMu              sync.RWMutex
+	subs                map[channelID]subWrapper
+	eventHandler        centrifuge.BrokerEventHandler
+	controlEventHandler centrifuge.ControlEventHandler
+	clientChannelPrefix string
+	rawModeReplacer     *strings.Replacer
+}
+
+var _ centrifuge.Broker = (*NatsBroker)(nil)
+var _ centrifuge.Controller = (*NatsBroker)(nil)
+
+// New creates NatsBroker.
+func New(n *centrifuge.Node, conf Config) (*NatsBroker, error) {
+	b := &NatsBroker{
+		node:                n,
+		config:              conf,
+		subs:                make(map[channelID]subWrapper),
+		clientChannelPrefix: conf.Prefix + ".client.",
+	}
+	if conf.RawMode.Enabled {
+		log.Info().Str("raw_mode_prefix", conf.RawMode.Prefix).Msg("raw mode of Nats enabled")
+		if len(conf.RawMode.ChannelReplacements) > 0 {
+			var replacerArgs []string
+			for k, v := range conf.RawMode.ChannelReplacements {
+				replacerArgs = append(replacerArgs, k, v)
+			}
+			b.rawModeReplacer = strings.NewReplacer(replacerArgs...)
+		}
+	}
+	return b, nil
+}
+
+func (b *NatsBroker) controlChannel() channelID {
+	return channelID(b.config.Prefix + ".control")
+}
+
+func (b *NatsBroker) nodeChannel(nodeID string) channelID {
+	return channelID(b.config.Prefix + ".node." + nodeID)
+}
+
+func (b *NatsBroker) clientChannel(ch string) channelID {
+	if b.config.RawMode.Enabled {
+		if b.rawModeReplacer != nil {
+			ch = b.rawModeReplacer.Replace(ch)
+		}
+		return channelID(b.config.RawMode.Prefix + ch)
+	}
+	return channelID(b.clientChannelPrefix + ch)
+}
+
+func (b *NatsBroker) initConnection() error {
+	var initErr error
+	b.initConnOnce.Do(func() {
+		url := b.config.URL
+		if url == "" {
+			url = nats.DefaultURL
+		}
+		options := []nats.Option{
+			nats.ReconnectBufSize(-1),
+			nats.MaxReconnects(-1),
+			nats.Timeout(b.config.DialTimeout.ToDuration()),
+			nats.FlusherTimeout(b.config.WriteTimeout.ToDuration()),
+		}
+		if b.config.TLS.Enabled {
+			tlsConfig, err := b.config.TLS.ToGoTLSConfig("nats")
+			if err != nil {
+				initErr = fmt.Errorf("error creating TLS config: %w", err)
+				return
+			}
+			options = append(options, nats.Secure(tlsConfig))
+		}
+		nc, err := nats.Connect(url, options...)
+		if err != nil {
+			initErr = fmt.Errorf("error connecting to %s: %w", url, err)
+			return
+		}
+		b.nc = nc
+	})
+	return initErr
+}
+
+func (b *NatsBroker) RegisterControlEventHandler(h centrifuge.ControlEventHandler) error {
+	b.controlEventHandler = h
+	err := b.initConnection()
+	if err != nil {
+		return fmt.Errorf("error initializing nats connection: %w", err)
+	}
+	_, err = b.nc.Subscribe(string(b.controlChannel()), b.handleControl)
+	if err != nil {
+		return err
+	}
+	_, err = b.nc.Subscribe(string(b.nodeChannel(b.node.ID())), b.handleControl)
+	if err != nil {
+		return err
+	}
+	return nil
+}
+
+func (b *NatsBroker) RegisterBrokerEventHandler(h centrifuge.BrokerEventHandler) error {
+	b.eventHandler = h
+	err := b.initConnection()
+	if err != nil {
+		return fmt.Errorf("error initializing nats connection: %w", err)
+	}
+	log.Info().Str("broker", "nats").Str("url", tools.StripPassword(b.config.URL)).Msg("broker running")
+	return nil
+}
+
+// Close ...
+func (b *NatsBroker) Close(_ context.Context) error {
+	b.nc.Close()
+	return nil
+}
+
+func (b *NatsBroker) IsSupportedSubscribeChannel(ch string) bool {
+	if b.config.AllowWildcards {
+		return true
+	}
+	if strings.Contains(ch, "*") || strings.Contains(ch, ">") {
+		return false
+	}
+	return true
+}
+
+func (b *NatsBroker) IsSupportedPublishChannel(ch string) bool {
+	if strings.Contains(ch, "*") || strings.Contains(ch, ">") {
+		return false
+	}
+	return true
+}
+
+// Publish - see Broker interface description.
+func (b *NatsBroker) Publish(ch string, data []byte, opts centrifuge.PublishOptions) (centrifuge.PublishResult, error) {
+	if !b.IsSupportedPublishChannel(ch) {
+		// Do not support wildcard subscriptions.
+		return centrifuge.PublishResult{}, centrifuge.ErrorBadRequest
+	}
+	if b.config.RawMode.Enabled {
+		return centrifuge.PublishResult{}, b.nc.Publish(string(b.clientChannel(ch)), data)
+	}
+	push := &protocol.Push{
+		Channel: ch,
+		Pub: &protocol.Publication{
+			Data:  data,
+			Info:  infoToProto(opts.ClientInfo),
+			Tags:  opts.Tags,
+			Delta: opts.UseDelta, // Will be cleaned up before passing to Node.
+			Time:  time.Now().UnixMilli(),
+		},
+	}
+	byteMessage, err := push.MarshalVT()
+	if err != nil {
+		return centrifuge.PublishResult{}, err
+	}
+	return centrifuge.PublishResult{}, b.nc.Publish(string(b.clientChannel(ch)), byteMessage)
+}
+
+// We experimentally support attaching epoch to publication.
+const epochTagsKey = "__centrifugo_epoch"
+
+func (b *NatsBroker) PublishWithStreamPosition(ch string, data []byte, opts centrifuge.PublishOptions, sp centrifuge.StreamPosition) error {
+	if !b.IsSupportedPublishChannel(ch) {
+		// Do not support wildcard subscriptions.
+		return centrifuge.ErrorBadRequest
+	}
+	if b.config.RawMode.Enabled {
+		// Do not support stream positions in raw mode.
+		return centrifuge.ErrorBadRequest
+	}
+	tags := opts.Tags
+	if tags == nil {
+		tags = map[string]string{}
+	}
+	tags[epochTagsKey] = sp.Epoch
+	push := &protocol.Push{
+		Channel: ch,
+		Pub: &protocol.Publication{
+			Offset: sp.Offset,
+			Data:   data,
+			Info:   infoToProto(opts.ClientInfo),
+			Tags:   tags,
+			Delta:  opts.UseDelta, // Will be cleaned up before passing to Node.
+		},
+	}
+	byteMessage, err := push.MarshalVT()
+	if err != nil {
+		return err
+	}
+	return b.nc.Publish(string(b.clientChannel(ch)), byteMessage)
+}
+
+// PublishJoin - see Broker interface description.
+func (b *NatsBroker) PublishJoin(ch string, info *centrifuge.ClientInfo) error {
+	if b.config.RawMode.Enabled {
+		return nil
+	}
+	push := &protocol.Push{
+		Channel: ch,
+		Join: &protocol.Join{
+			Info: infoToProto(info),
+		},
+	}
+	byteMessage, err := push.MarshalVT()
+	if err != nil {
+		return err
+	}
+	return b.nc.Publish(string(b.clientChannel(ch)), byteMessage)
+}
+
+// PublishLeave - see Broker interface description.
+func (b *NatsBroker) PublishLeave(ch string, info *centrifuge.ClientInfo) error {
+	if b.config.RawMode.Enabled {
+		return nil
+	}
+	push := &protocol.Push{
+		Channel: ch,
+		Leave: &protocol.Leave{
+			Info: infoToProto(info),
+		},
+	}
+	byteMessage, err := push.MarshalVT()
+	if err != nil {
+		return err
+	}
+	return b.nc.Publish(string(b.clientChannel(ch)), byteMessage)
+}
+
+// PublishControl - see Broker interface description.
+func (b *NatsBroker) PublishControl(data []byte, nodeID, _ string) error {
+	var channelID channelID
+	if nodeID == "" {
+		channelID = b.controlChannel()
+	} else {
+		channelID = b.nodeChannel(nodeID)
+	}
+	return b.nc.Publish(string(channelID), data)
+}
+
+// History ...
+func (b *NatsBroker) History(_ string, _ centrifuge.HistoryOptions) ([]*centrifuge.Publication, centrifuge.StreamPosition, error) {
+	return nil, centrifuge.StreamPosition{}, centrifuge.ErrorNotAvailable
+}
+
+// RemoveHistory ...
+func (b *NatsBroker) RemoveHistory(_ string) error {
+	return centrifuge.ErrorNotAvailable
+}
+
+func (b *NatsBroker) handleClientMessage(subject string, data []byte, sub *nats.Subscription) {
+	if b.config.RawMode.Enabled {
+		b.subsMu.RLock()
+		subWrap, ok := b.subs[channelID(sub.Subject)]
+		b.subsMu.RUnlock()
+		if !ok {
+			return
+		}
+		channel := subWrap.origChannel
+		_ = b.eventHandler.HandlePublication(
+			channel,
+			&centrifuge.Publication{Data: data, Channel: strings.TrimPrefix(subject, b.config.RawMode.Prefix)},
+			centrifuge.StreamPosition{}, false, nil)
+		return
+	}
+
+	var push protocol.Push
+	err := push.UnmarshalVT(data)
+	if err != nil {
+		log.Error().Err(err).Msg("error unmarshal push from Nats")
+		return
+	}
+
+	if push.Pub != nil {
+		var subChannel = push.Channel
+		var specificChannel string
+		if b.config.AllowWildcards {
+			subChannel = strings.TrimPrefix(sub.Subject, b.clientChannelPrefix)
+			specificChannel = push.Channel
+		}
+		sp := centrifuge.StreamPosition{}
+		if push.Pub.Offset > 0 && push.Pub.Tags != nil {
+			sp.Offset = push.Pub.Offset
+			sp.Epoch = push.Pub.Tags[epochTagsKey]
+		}
+		// Strip the internal epoch tag so it isn't delivered to clients as if it
+		// were an application tag (no-op when absent).
+		delete(push.Pub.Tags, epochTagsKey)
+		delta := push.Pub.Delta
+		push.Pub.Delta = false
+		_ = b.eventHandler.HandlePublication(subChannel, pubFromProto(push.Pub, specificChannel), sp, delta, nil)
+	} else if push.Join != nil {
+		_ = b.eventHandler.HandleJoin(push.Channel, infoFromProto(push.Join.Info))
+	} else if push.Leave != nil {
+		_ = b.eventHandler.HandleLeave(push.Channel, infoFromProto(push.Leave.Info))
+	} else {
+		log.Warn().Str("push", fmt.Sprintf("%v", &push)).Msg("unknown push from Nats")
+	}
+}
+
+func (b *NatsBroker) handleClient(m *nats.Msg) {
+	b.handleClientMessage(m.Subject, m.Data, m.Sub)
+}
+
+func (b *NatsBroker) handleControl(m *nats.Msg) {
+	_ = b.controlEventHandler.HandleControl(m.Data)
+}
+
+// Subscribe - see Broker interface description.
+func (b *NatsBroker) Subscribe(channels ...string) error {
+	for _, ch := range channels {
+		if !b.IsSupportedSubscribeChannel(ch) {
+			// Do not support wildcard subscriptions.
+			return centrifuge.ErrorBadRequest
+		}
+		clientChannel := b.clientChannel(ch)
+		b.subsMu.RLock()
+		_, ok := b.subs[clientChannel]
+		b.subsMu.RUnlock()
+		if ok {
+			continue
+		}
+		subscription, err := b.nc.Subscribe(string(clientChannel), b.handleClient)
+		if err != nil {
+			return err
+		}
+		b.subsMu.Lock()
+		b.subs[clientChannel] = subWrapper{
+			sub:         subscription,
+			origChannel: ch,
+		}
+		b.subsMu.Unlock()
+	}
+	return nil
+}
+
+// Unsubscribe - see Broker interface description.
+func (b *NatsBroker) Unsubscribe(channels ...string) error {
+	for _, ch := range channels {
+		clientChannel := b.clientChannel(ch)
+		b.subsMu.RLock()
+		subWrap, ok := b.subs[clientChannel]
+		b.subsMu.RUnlock()
+		if ok {
+			err := subWrap.sub.Unsubscribe()
+			if err != nil {
+				return err
+			}
+			b.subsMu.Lock()
+			delete(b.subs, clientChannel)
+			b.subsMu.Unlock()
+		}
+	}
+	return nil
+}
+
+func infoFromProto(v *protocol.ClientInfo) *centrifuge.ClientInfo {
+	if v == nil {
+		return nil
+	}
+	info := &centrifuge.ClientInfo{
+		ClientID: v.GetClient(),
+		UserID:   v.GetUser(),
+	}
+	if len(v.ConnInfo) > 0 {
+		info.ConnInfo = v.ConnInfo
+	}
+	if len(v.ChanInfo) > 0 {
+		info.ChanInfo = v.ChanInfo
+	}
+	return info
+}
+
+func infoToProto(v *centrifuge.ClientInfo) *protocol.ClientInfo {
+	if v == nil {
+		return nil
+	}
+	info := &protocol.ClientInfo{
+		Client: v.ClientID,
+		User:   v.UserID,
+	}
+	if len(v.ConnInfo) > 0 {
+		info.ConnInfo = v.ConnInfo
+	}
+	if len(v.ChanInfo) > 0 {
+		info.ChanInfo = v.ChanInfo
+	}
+	return info
+}
+
+func pubFromProto(pub *protocol.Publication, specificChannel string) *centrifuge.Publication {
+	if pub == nil {
+		return nil
+	}
+	return &centrifuge.Publication{
+		Offset:  pub.GetOffset(),
+		Data:    pub.Data,
+		Info:    infoFromProto(pub.GetInfo()),
+		Channel: specificChannel,
+		Tags:    pub.GetTags(),
+		Time:    pub.Time,
+	}
+}

internal/natsbroker/broker_test.go

@@ -0,0 +1,326 @@
+//go:build integration
+
+package natsbroker
+
+import (
+	"context"
+	"math/rand"
+	"strconv"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/centrifugal/centrifugo/v6/internal/configtypes"
+
+	"github.com/centrifugal/centrifuge"
+	"github.com/stretchr/testify/require"
+)
+
+func newTestNatsBroker() *NatsBroker {
+	return NewTestNatsBrokerWithPrefix("centrifugo-test")
+}
+
+func NewTestNatsBrokerWithPrefix(prefix string) *NatsBroker {
+	n, _ := centrifuge.New(centrifuge.Config{})
+	b, _ := New(n, Config{NatsPrefixed: configtypes.NatsPrefixed{
+		Prefix: prefix,
+	}})
+	n.SetBroker(b)
+	err := n.Run()
+	if err != nil {
+		panic(err)
+	}
+	return b
+}
+
+func BenchmarkNatsEnginePublish(b *testing.B) {
+	broker := newTestNatsBroker()
+	rawData := []byte(`{"bench": true}`)
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		_, err := broker.Publish("channel", rawData, centrifuge.PublishOptions{})
+		if err != nil {
+			panic(err)
+		}
+	}
+}
+
+func BenchmarkNatsEnginePublishParallel(b *testing.B) {
+	broker := newTestNatsBroker()
+	rawData := []byte(`{"bench": true}`)
+	b.SetParallelism(128)
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			_, err := broker.Publish("channel", rawData, centrifuge.PublishOptions{})
+			if err != nil {
+				panic(err)
+			}
+		}
+	})
+}
+
+func BenchmarkNatsEngineSubscribe(b *testing.B) {
+	broker := newTestNatsBroker()
+	j := 0
+	b.ResetTimer()
+	for i := 0; i < b.N; i++ {
+		j++
+		err := broker.Subscribe("subscribe" + strconv.Itoa(j))
+		if err != nil {
+			panic(err)
+		}
+	}
+}
+
+func BenchmarkNatsEngineSubscribeParallel(b *testing.B) {
+	broker := newTestNatsBroker()
+	i := 0
+	b.SetParallelism(128)
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			i++
+			err := broker.Subscribe("subscribe" + strconv.Itoa(i))
+			if err != nil {
+				panic(err)
+			}
+		}
+	})
+}
+
+type natsTest struct {
+	Name         string
+	BrokerConfig Config
+	// Not all configurations support join/leave messages.
+	TestJoinLeave bool
+	// For WildcardChannel case we subscribe once to a wildcard channel instead of individual channels.
+	WildcardChannel bool
+}
+
+var natsTests = []natsTest{
+	{"default_mode", Config{}, true, false},
+	{"raw_mode", Config{RawMode: configtypes.RawModeConfig{Enabled: true}}, false, false},
+	{"raw_mode_wildcards", Config{AllowWildcards: true, RawMode: configtypes.RawModeConfig{Enabled: true}}, false, true},
+}
+
+var letterRunes = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
+
+func randString(n int) string {
+	random := rand.New(rand.NewSource(time.Now().UnixNano()))
+	b := make([]rune, n)
+	for i := range b {
+		b[i] = letterRunes[random.Intn(len(letterRunes))]
+	}
+	return string(b)
+}
+
+func getUniquePrefix() string {
+	return "centrifugo-test-" + randString(3) + "-" + strconv.FormatInt(time.Now().UnixNano(), 10)
+}
+
+func stopNatsBroker(b *NatsBroker) {
+	_ = b.Close(context.Background())
+}
+
+func TestNatsPubSubTwoNodes(t *testing.T) {
+	for _, tt := range natsTests {
+		t.Run(tt.Name, func(t *testing.T) {
+			prefix := getUniquePrefix()
+			tt.BrokerConfig.Prefix = prefix
+			tt.BrokerConfig.RawMode.Prefix = prefix
+
+			node1, err := centrifuge.New(centrifuge.Config{})
+			require.NoError(t, err)
+			b1, _ := New(node1, tt.BrokerConfig)
+			node1.SetBroker(b1)
+			defer func() { _ = node1.Shutdown(context.Background()) }()
+			defer stopNatsBroker(b1)
+
+			msgNum := 10
+			var numPublications int64
+			var numJoins int64
+			var numLeaves int64
+			pubCh := make(chan struct{})
+			joinCh := make(chan struct{})
+			leaveCh := make(chan struct{})
+			brokerEventHandler := &testBrokerEventHandler{
+				HandleControlFunc: func(bytes []byte) error {
+					return nil
+				},
+				HandlePublicationFunc: func(ch string, pub *centrifuge.Publication, sp centrifuge.StreamPosition, delta bool, prevPub *centrifuge.Publication) error {
+					c := atomic.AddInt64(&numPublications, 1)
+					if c == int64(msgNum) {
+						close(pubCh)
+					}
+					return nil
+				},
+				HandleJoinFunc: func(ch string, info *centrifuge.ClientInfo) error {
+					c := atomic.AddInt64(&numJoins, 1)
+					if c == int64(msgNum) {
+						close(joinCh)
+					}
+					return nil
+				},
+				HandleLeaveFunc: func(ch string, info *centrifuge.ClientInfo) error {
+					c := atomic.AddInt64(&numLeaves, 1)
+					if c == int64(msgNum) {
+						close(leaveCh)
+					}
+					return nil
+				},
+			}
+			_ = b1.RegisterControlEventHandler(brokerEventHandler)
+			_ = b1.RegisterBrokerEventHandler(brokerEventHandler)
+
+			if tt.WildcardChannel {
+				require.NoError(t, b1.Subscribe("test.*"))
+			} else {
+				for i := 0; i < msgNum; i++ {
+					require.NoError(t, b1.Subscribe("test."+strconv.Itoa(i)))
+				}
+			}
+
+			node2, _ := centrifuge.New(centrifuge.Config{})
+
+			b2, _ := New(node2, tt.BrokerConfig)
+			node2.SetBroker(b2)
+			_ = node2.Run()
+			defer func() { _ = node2.Shutdown(context.Background()) }()
+			defer stopNatsBroker(b2)
+
+			for i := 0; i < msgNum; i++ {
+				_, err := node2.Publish("test."+strconv.Itoa(i), []byte("123"))
+				require.NoError(t, err)
+				if tt.TestJoinLeave {
+					err = b2.PublishJoin("test."+strconv.Itoa(i), &centrifuge.ClientInfo{})
+					require.NoError(t, err)
+					err = b2.PublishLeave("test."+strconv.Itoa(i), &centrifuge.ClientInfo{})
+					require.NoError(t, err)
+				}
+			}
+
+			select {
+			case <-pubCh:
+			case <-time.After(time.Second):
+				require.Fail(t, "timeout waiting for PUB/SUB message")
+			}
+			if tt.TestJoinLeave {
+				select {
+				case <-joinCh:
+				case <-time.After(time.Second):
+					require.Fail(t, "timeout waiting for PUB/SUB join message")
+				}
+				select {
+				case <-leaveCh:
+				case <-time.After(time.Second):
+					require.Fail(t, "timeout waiting for PUB/SUB leave message")
+				}
+			}
+		})
+	}
+}
+
+func TestNatsPubSubTwoNodes_ChannelReplacement(t *testing.T) {
+	brokerConfig := Config{RawMode: configtypes.RawModeConfig{
+		Enabled: true,
+	}}
+	prefix := getUniquePrefix()
+	brokerConfig.Prefix = prefix
+	brokerConfig.RawMode.Prefix = prefix
+	brokerConfig.RawMode.ChannelReplacements = map[string]string{
+		":": ".",
+	}
+	node1, err := centrifuge.New(centrifuge.Config{})
+	require.NoError(t, err)
+	b1, _ := New(node1, brokerConfig)
+	node1.SetBroker(b1)
+	defer func() { _ = node1.Shutdown(context.Background()) }()
+	defer stopNatsBroker(b1)
+
+	msgNum := 10
+	var numPublications int64
+	pubCh := make(chan struct{})
+	brokerEventHandler := &testBrokerEventHandler{
+		HandleControlFunc: func(bytes []byte) error {
+			return nil
+		},
+		HandlePublicationFunc: func(ch string, pub *centrifuge.Publication, sp centrifuge.StreamPosition, delta bool, prevPub *centrifuge.Publication) error {
+			c := atomic.AddInt64(&numPublications, 1)
+			if c == int64(msgNum) {
+				close(pubCh)
+			}
+			return nil
+		},
+		HandleJoinFunc: func(ch string, info *centrifuge.ClientInfo) error {
+			return nil
+		},
+		HandleLeaveFunc: func(ch string, info *centrifuge.ClientInfo) error {
+			return nil
+		},
+	}
+	_ = b1.RegisterControlEventHandler(brokerEventHandler)
+	_ = b1.RegisterBrokerEventHandler(brokerEventHandler)
+
+	for i := 0; i < msgNum; i++ {
+		require.NoError(t, b1.Subscribe("test:"+strconv.Itoa(i)))
+	}
+
+	node2, _ := centrifuge.New(centrifuge.Config{})
+
+	b2, _ := New(node2, brokerConfig)
+	node2.SetBroker(b2)
+	_ = node2.Run()
+	defer func() { _ = node2.Shutdown(context.Background()) }()
+	defer stopNatsBroker(b2)
+
+	for i := 0; i < msgNum; i++ {
+		_, err := node2.Publish("test:"+strconv.Itoa(i), []byte("123"))
+		require.NoError(t, err)
+	}
+
+	select {
+	case <-pubCh:
+	case <-time.After(3 * time.Second):
+		require.Fail(t, "timeout waiting for PUB/SUB message")
+	}
+}
+
+type testBrokerEventHandler struct {
+	// Publication must register callback func to handle Publications received.
+	HandlePublicationFunc func(ch string, pub *centrifuge.Publication, sp centrifuge.StreamPosition, delta bool, prevPub *centrifuge.Publication) error
+	// Join must register callback func to handle Join messages received.
+	HandleJoinFunc func(ch string, info *centrifuge.ClientInfo) error
+	// Leave must register callback func to handle Leave messages received.
+	HandleLeaveFunc func(ch string, info *centrifuge.ClientInfo) error
+	// Control must register callback func to handle Control data received.
+	HandleControlFunc func([]byte) error
+}
+
+func (b *testBrokerEventHandler) HandlePublication(ch string, pub *centrifuge.Publication, sp centrifuge.StreamPosition, delta bool, prevPub *centrifuge.Publication) error {
+	if b.HandlePublicationFunc != nil {
+		return b.HandlePublicationFunc(ch, pub, sp, delta, prevPub)
+	}
+	return nil
+}
+
+func (b *testBrokerEventHandler) HandleJoin(ch string, info *centrifuge.ClientInfo) error {
+	if b.HandleJoinFunc != nil {
+		return b.HandleJoinFunc(ch, info)
+	}
+	return nil
+}
+
+func (b *testBrokerEventHandler) HandleLeave(ch string, info *centrifuge.ClientInfo) error {
+	if b.HandleLeaveFunc != nil {
+		return b.HandleLeaveFunc(ch, info)
+	}
+	return nil
+}
+
+func (b *testBrokerEventHandler) HandleControl(data []byte) error {
+	if b.HandleControlFunc != nil {
+		return b.HandleControlFunc(data)
+	}
+	return nil
+}