diff --git a/waku/v2/api/publish/message_sender.go b/waku/v2/api/publish/message_sender.go index 1cb60a604..7ccfe8404 100644 --- a/waku/v2/api/publish/message_sender.go +++ b/waku/v2/api/publish/message_sender.go @@ -79,14 +79,19 @@ func (r *Request) WithPublishMethod(publishMethod PublishMethod) *Request { return r } -func NewMessageSender(publishMethod PublishMethod, publisher Publisher, logger *zap.Logger) (*MessageSender, error) { +func NewMessageSender(publishMethod PublishMethod, publisher Publisher, rateLimiter PublishRateLimiter, logger *zap.Logger) (*MessageSender, error) { if publishMethod == UnknownMethod { return nil, errors.New("publish method is required") } + + if rateLimiter == nil { + rateLimiter = NewDefaultRateLimiter(DefaultPublishingLimiterRate, DefaultPublishingLimitBurst) + } + return &MessageSender{ publishMethod: publishMethod, publisher: publisher, - rateLimiter: NewRlnRateLimiter(RlnLimiterCapacity, RlnLimiterRefillInterval), + rateLimiter: rateLimiter, logger: logger, }, nil } diff --git a/waku/v2/api/publish/message_sender_test.go b/waku/v2/api/publish/message_sender_test.go index c770992c0..7266b65fd 100644 --- a/waku/v2/api/publish/message_sender_test.go +++ b/waku/v2/api/publish/message_sender_test.go @@ -41,7 +41,7 @@ func (m *MockMessageSentCheck) Start() { } func TestNewSenderWithUnknownMethod(t *testing.T) { - sender, err := NewMessageSender(UnknownMethod, nil, nil) + sender, err := NewMessageSender(UnknownMethod, nil, nil, nil) require.NotNil(t, err) require.Nil(t, sender) } @@ -55,7 +55,7 @@ func TestNewSenderWithRelay(t *testing.T) { _, err = relayNode.Subscribe(context.Background(), protocol.NewContentFilter("test-pubsub-topic")) require.Nil(t, err) publisher := NewDefaultPublisher(nil, relayNode) - sender, err := NewMessageSender(Relay, publisher, utils.Logger()) + sender, err := NewMessageSender(Relay, publisher, nil, utils.Logger()) require.Nil(t, err) require.NotNil(t, sender) require.Nil(t, sender.messageSentCheck) @@ -81,7 +81,7 @@ func TestNewSenderWithRelayAndMessageSentCheck(t *testing.T) { _, err = relayNode.Subscribe(context.Background(), protocol.NewContentFilter("test-pubsub-topic")) require.Nil(t, err) publisher := NewDefaultPublisher(nil, relayNode) - sender, err := NewMessageSender(Relay, publisher, utils.Logger()) + sender, err := NewMessageSender(Relay, publisher, nil, utils.Logger()) check := &MockMessageSentCheck{Messages: make(map[string]map[common.Hash]uint32)} sender.WithMessageSentCheck(check) @@ -111,7 +111,7 @@ func TestNewSenderWithRelayAndMessageSentCheck(t *testing.T) { } func TestNewSenderWithLightPush(t *testing.T) { - sender, err := NewMessageSender(LightPush, nil, nil) + sender, err := NewMessageSender(LightPush, nil, nil, nil) require.Nil(t, err) require.NotNil(t, sender) require.Equal(t, LightPush, sender.publishMethod) @@ -140,7 +140,7 @@ func TestMessageSentEmitter(t *testing.T) { _, err = relayNode.Subscribe(context.Background(), protocol.NewContentFilter("test-pubsub-topic")) require.Nil(t, err) publisher := NewDefaultPublisher(nil, relayNode) - sender, err := NewMessageSender(Relay, publisher, utils.Logger()) + sender, err := NewMessageSender(Relay, publisher, nil, utils.Logger()) require.Nil(t, err) check := &MockMessageSentCheck{Messages: make(map[string]map[common.Hash]uint32)}