Repository: apache/pulsar-client-go
Pull request #955 (2023-02-15 13:59:46)

[feat] Support memory limit for producer.
Master Issue: #927 

### Motivation

#927 

### Modifications

- Added `channel_cond` class to enhance `sync.cond` to support wait that can accept context.
- Added the `memory_limit_controller` class.
- Producer adds memory limit relate logic.


### Verifying this change
- Added `channel_cond_test` to cover `channel_cond`
- Added `memory_limit_controller_test` to cover `mem_mory_limit_controller`
- Added `TestMemLimitRejectProducerMessages` and `TestMemLimitContextCancel ` to cover producer relate logic.

### Does this pull request potentially affect one of the following parts:

*If `yes` was chosen, please highlight the changes*

  - Dependencies (does it add or upgrade a dependency): (no)
  - The public API: (yes)
  - The schema: (no)
  - The default values of configurations: (no)
  - The wire protocol: (no)

### Documentation
  - Does this pull request introduce a new feature? (yes)
  - If yes, how is the feature documented? (GoDocs)
  


shibd: > And I don't get why do we have to wrap a conditional variable? 

The native cond does not support timeout wait, and mainly wants to control the timeout or cancel the wait through context.

https://github.com/apache/pulsar-client-go/pull/955/files#diff-d4362cb0b0046577b68820bba2bde37b97ca5d7db7237e4e4b5eac54d7aa8bc3R89-R120

> We can achieve the same goal with the channel itself,

Do you want `mem_controller_limit` all the methods into the channel first, and then select to read the signal and process it? Similar to #957 `ackGroupingTracker `? 

> or just reuse the Semaphore [here](https://github.com/apache/pulsar-client-go/blob/master/pulsar/internal/semaphore.go)

This semaphore can only process applications or release 1 signal.



BewareMyPower: For example, if we add a `ch chan bool` field to `memoryLimitController`, we can implement blocking `ReserveMemory` method like:

```go
func (m *memoryLimitController) ReserveMemory(ctx context.Context, size int64) bool {
        currentUsage := atomic.AddInt64(&m.currentUsage, size)
        for currentUsage > m.limit {
                select {
                case <-m.ch:
                        currentUsage = atomic.LoadInt64(&m.currentUsage)
                case <-ctx.Done():
                        return false
                }
        }
        return true
}

func (m *memoryLimitController) ReleaseMemory(size int64) {
        newUsage := atomic.AddInt64(&m.currentUsage, -size)
        // newUsage+size > m.limit means m was blocked in ReserveMemory method
        if newUsage+size > m.limit && newUsage <= m.limit {
                m.ch <- true
        }
}
```

The code above is not verified yet. But with the channel the code looks more simple and clear.
shibd: > > This semaphore can only support acquire or release 1 signal.
> 
> Yes. But we can adopt a similar way to handle this case.
> 
> For example, if we add a `ch chan bool` field to `memoryLimitController`, we can implement blocking `ReserveMemory` method like:
> 
> ```go
> func (m *memoryLimitController) ReserveMemory(ctx context.Context, size int64) bool {
>         // NOTE: maybe we need to check if m.currentUsage > m.limit first
>         currentUsage := atomic.AddInt64(&m.currentUsage, size)
>         for currentUsage > m.limit {
>                 select {
>                 case <-m.ch:
>                         currentUsage = atomic.LoadInt64(&m.currentUsage)
>                 case <-ctx.Done(): // NOTE: Not sure if we need to reset some fields here
>                         return false
>                 }
>         }
>         return true
> }
> 
> func (m *memoryLimitController) ReleaseMemory(size int64) {
>         newUsage := atomic.AddInt64(&m.currentUsage, -size)
>         // newUsage+size > m.limit means m was blocked in ReserveMemory method
>         if newUsage+size > m.limit && newUsage <= m.limit {
>                 m.ch <- true
>         }
> }
> ```
> 
> The code above is not verified yet. But with the channel the code looks more simple and clear.

This implementation has a problem, and the `broadcasting` cannot be implemented using channels.

For example:
1. `currentUsage` = 100, `limit` = 100
2. goroutine `1` call `ReseveMemory(10)`, it will blocked.
3. goroutine `2` call `ReseveMemory(10)`, it will blocked.
4. goroutine `3` call `ReleaseMemory(20)`, it only wake one goroutine(1 or 2). The expectation is that both are woken up and return true.


// NOTE: maybe we need to check if m.currentUsage > m.limit first

Like this note. If we want to handle this case, we may need to introduce a variable `waitNum`. So I feel like it's complicated to implement `mem_controller_limit`, and it's better to introduce `channel_cond`.


BewareMyPower: > If we want to handle this case, we may need to introduce a variable waitNum

I see. Currently let's go head with the conditional variable, but don't expose this struct to users.
RobertIndie: > Like this note. If we want to handle this case, we may need to introduce a variable waitNum. So I feel like it's complicated to implement mem_controller_limit, and it's better to introduce channel_cond.

@shibd For this case, another workaround is to close the channel instead of sending value. But it's actually what this PR does.  


Patch:
diff --git a/pulsar/client.go b/pulsar/client.go
index 135d22b..75b363d 100644
--- a/pulsar/client.go
+++ b/pulsar/client.go
@@ -144,6 +144,10 @@ type ClientOptions struct {
 	MetricsRegisterer prometheus.Registerer
 
 	EnableTransaction bool
+
+	// Limit of client memory usage (in byte). The 64M default can guarantee a high producer throughput.
+	// Config less than 0 indicates off memory limit.
+	MemoryLimitBytes int64
 }
 
 // Client represents a pulsar client
diff --git a/pulsar/client_impl.go b/pulsar/client_impl.go
index f444804..7d90922 100644
--- a/pulsar/client_impl.go
+++ b/pulsar/client_impl.go
@@ -33,6 +33,7 @@ const (
 	defaultConnectionTimeout = 10 * time.Second
 	defaultOperationTimeout  = 30 * time.Second
 	defaultKeepAliveInterval = 30 * time.Second
+	defaultMemoryLimitBytes  = 64 * 1024 * 1024
 )
 
 type client struct {
@@ -42,6 +43,7 @@ type client struct {
 	lookupService internal.LookupService
 	metrics       *internal.Metrics
 	tcClient      *transactionCoordinatorClient
+	memLimit      internal.MemoryLimitController
 
 	log log.Logger
 }
@@ -134,11 +136,17 @@ func newClient(options ClientOptions) (Client, error) {
 		keepAliveInterval = defaultKeepAliveInterval
 	}
 
+	memLimitBytes := options.MemoryLimitBytes
+	if memLimitBytes == 0 {
+		memLimitBytes = defaultMemoryLimitBytes
+	}
+
 	c := &client{
 		cnxPool: internal.NewConnectionPool(tlsConfig, authProvider, connectionTimeout, keepAliveInterval,
 			maxConnectionsPerHost, logger, metrics),
-		log:     logger,
-		metrics: metrics,
+		log:      logger,
+		metrics:  metrics,
+		memLimit: internal.NewMemoryLimitController(memLimitBytes),
 	}
 	serviceNameResolver := internal.NewPulsarServiceNameResolver(url)
 
diff --git a/pulsar/error.go b/pulsar/error.go
index ce366f5..0aa1e3c 100644
--- a/pulsar/error.go
+++ b/pulsar/error.go
@@ -110,6 +110,9 @@ const (
 	InvalidStatus
 	// TransactionError means this is a transaction related error
 	TransactionError
+
+	// ClientMemoryBufferIsFull client limit buffer is full
+	ClientMemoryBufferIsFull
 )
 
 // Error implement error interface, composed of two parts: msg and result.
@@ -216,6 +219,8 @@ func getResultStr(r Result) string {
 		return "ProducerClosed"
 	case SchemaFailure:
 		return "SchemaFailure"
+	case ClientMemoryBufferIsFull:
+		return "ClientMemoryBufferIsFull"
 	default:
 		return fmt.Sprintf("Result(%d)", r)
 	}
diff --git a/pulsar/internal/channel_cond.go b/pulsar/internal/channel_cond.go
new file mode 100644
index 0000000..38301ab
--- /dev/null
+++ b/pulsar/internal/channel_cond.go
@@ -0,0 +1,76 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package internal
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+	"unsafe"
+)
+
+type chCond struct {
+	L sync.Locker
+	// The pointer to the channel, the channel pointed to may change,
+	// because we will use the channel's close mechanism to implement broadcast notifications.
+	notifyChPtr unsafe.Pointer
+}
+
+func newCond(l sync.Locker) *chCond {
+	c := &chCond{L: l}
+	n := make(chan struct{})
+	c.notifyChPtr = unsafe.Pointer(&n)
+	return c
+}
+
+// wait for broadcast calls. Similar to regular sync.Cond
+func (c *chCond) wait() {
+	n := c.notifyChan()
+	c.L.Unlock()
+	<-n
+	c.L.Lock()
+}
+
+// waitWithContext Same as wait() call, but the end condition can also be controlled through the context.
+func (c *chCond) waitWithContext(ctx context.Context) bool {
+	n := c.notifyChan()
+	c.L.Unlock()
+	defer c.L.Lock()
+	select {
+	case <-n:
+		return true
+	case <-ctx.Done():
+		return false
+	default:
+		return true
+	}
+}
+
+// broadcast wakes all goroutines waiting on c.
+// It is not required for the caller to hold c.L during the call.
+func (c *chCond) broadcast() {
+	n := make(chan struct{})
+	ptrOld := atomic.SwapPointer(&c.notifyChPtr, unsafe.Pointer(&n))
+	// close old channels to trigger broadcast.
+	close(*(*chan struct{})(ptrOld))
+}
+
+func (c *chCond) notifyChan() <-chan struct{} {
+	ptr := atomic.LoadPointer(&c.notifyChPtr)
+	return *((*chan struct{})(ptr))
+}
diff --git a/pulsar/internal/memory_limit_controller.go b/pulsar/internal/memory_limit_controller.go
new file mode 100644
index 0000000..5bf8d59
--- /dev/null
+++ b/pulsar/internal/memory_limit_controller.go
@@ -0,0 +1,101 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package internal
+
+import (
+	"context"
+	"sync"
+	"sync/atomic"
+)
+
+type MemoryLimitController interface {
+	ReserveMemory(ctx context.Context, size int64) bool
+	TryReserveMemory(size int64) bool
+	ForceReserveMemory(size int64)
+	ReleaseMemory(size int64)
+	CurrentUsage() int64
+	CurrentUsagePercent() float64
+	IsMemoryLimited() bool
+}
+
+type memoryLimitController struct {
+	limit        int64
+	chCond       *chCond
+	currentUsage int64
+}
+
+func NewMemoryLimitController(limit int64) MemoryLimitController {
+	mlc := &memoryLimitController{
+		limit:  limit,
+		chCond: newCond(&sync.Mutex{}),
+	}
+	return mlc
+}
+
+func (m *memoryLimitController) ReserveMemory(ctx context.Context, size int64) bool {
+	if !m.TryReserveMemory(size) {
+		m.chCond.L.Lock()
+		defer m.chCond.L.Unlock()
+
+		for !m.TryReserveMemory(size) {
+			if !m.chCond.waitWithContext(ctx) {
+				return false
+			}
+		}
+	}
+	return true
+}
+
+func (m *memoryLimitController) TryReserveMemory(size int64) bool {
+	for {
+		current := atomic.LoadInt64(&m.currentUsage)
+		newUsage := current + size
+
+		// This condition means we allowed one request to go over the limit.
+		if m.IsMemoryLimited() && current > m.limit {
+			return false
+		}
+
+		if atomic.CompareAndSwapInt64(&m.currentUsage, current, newUsage) {
+			return true
+		}
+	}
+}
+
+func (m *memoryLimitController) ForceReserveMemory(size int64) {
+	atomic.AddInt64(&m.currentUsage, size)
+}
+
+func (m *memoryLimitController) ReleaseMemory(size int64) {
+	newUsage := atomic.AddInt64(&m.currentUsage, -size)
+	if newUsage+size > m.limit && newUsage <= m.limit {
+		m.chCond.broadcast()
+	}
+}
+
+func (m *memoryLimitController) CurrentUsage() int64 {
+	return atomic.LoadInt64(&m.currentUsage)
+}
+
+func (m *memoryLimitController) CurrentUsagePercent() float64 {
+	return float64(atomic.LoadInt64(&m.currentUsage)) / float64(m.limit)
+}
+
+func (m *memoryLimitController) IsMemoryLimited() bool {
+	return m.limit > 0
+}
diff --git a/pulsar/producer_partition.go b/pulsar/producer_partition.go
index 160693c..c3a0aa9 100644
--- a/pulsar/producer_partition.go
+++ b/pulsar/producer_partition.go
@@ -51,13 +51,14 @@ const (
 )
 
 var (
-	errFailAddToBatch  = newError(AddToBatchFailed, "message add to batch failed")
-	errSendTimeout     = newError(TimeoutError, "message send timeout")
-	errSendQueueIsFull = newError(ProducerQueueIsFull, "producer send queue is full")
-	errContextExpired  = newError(TimeoutError, "message send context expired")
-	errMessageTooLarge = newError(MessageTooBig, "message size exceeds MaxMessageSize")
-	errMetaTooLarge    = newError(InvalidMessage, "message metadata size exceeds MaxMessageSize")
-	errProducerClosed  = newError(ProducerClosed, "producer already been closed")
+	errFailAddToBatch     = newError(AddToBatchFailed, "message add to batch failed")
+	errSendTimeout        = newError(TimeoutError, "message send timeout")
+	errSendQueueIsFull    = newError(ProducerQueueIsFull, "producer send queue is full")
+	errContextExpired     = newError(TimeoutError, "message send context expired")
+	errMessageTooLarge    = newError(MessageTooBig, "message size exceeds MaxMessageSize")
+	errMetaTooLarge       = newError(InvalidMessage, "message metadata size exceeds MaxMessageSize")
+	errProducerClosed     = newError(ProducerClosed, "producer already been closed")
+	errMemoryBufferIsFull = newError(ClientMemoryBufferIsFull, "client memory buffer is full")
 
 	buffersPool sync.Pool
 )
@@ -483,6 +484,7 @@ func (p *partitionProducer) internalSend(request *sendRequest) {
 
 	// read payload from message
 	uncompressedPayload := msg.Payload
+	uncompressedPayloadSize := int64(len(uncompressedPayload))
 
 	var schemaPayload []byte
 	var err error
@@ -494,14 +496,14 @@ func (p *partitionProducer) internalSend(request *sendRequest) {
 
 	// The block chan must be closed when returned with exception
 	defer request.stopBlock()
-	if !p.canAddToQueue(request) {
+	if !p.canAddToQueue(request, uncompressedPayloadSize) {
 		return
 	}
 
 	if p.options.DisableMultiSchema {
 		if msg.Schema != nil && p.options.Schema != nil &&
 			msg.Schema.GetSchemaInfo().hash() != p.options.Schema.GetSchemaInfo().hash() {
-			p.publishSemaphore.Release()
+			p.releaseSemaphoreAndMem(uncompressedPayloadSize)
 			request.callback(nil, request.msg, fmt.Errorf("msg schema can not match with producer schema"))
 			p.log.WithError(err).Errorf("The producer %s of the topic %s is disabled the `MultiSchema`", p.producerName, p.topic)
 			return
@@ -520,7 +522,7 @@ func (p *partitionProducer) internalSend(request *sendRequest) {
 		if uncompressedPayload == nil && schema != nil {
 			schemaPayload, err = schema.Encode(msg.Value)
 			if err != nil {
-				p.publishSemaphore.Release()
+				p.releaseSemaphoreAndMem(uncompressedPayloadSize)
 				request.callback(nil, request.msg, newError(SchemaFailure, err.Error()))
 				p.log.WithError(err).Errorf("Schema encode message failed %s", msg.Value)
 				return
@@ -536,7 +538,7 @@ func (p *partitionProducer) internalSend(request *sendRequest) {
 		if schemaVersion == nil {
 			schemaVersion, err = p.getOrCreateSchema(schema.GetSchemaInfo())
 			if err != nil {
-				p.publishSemaphore.Release()
+				p.releaseSemaphoreAndMem(uncompressedPayloadSize)
 				p.log.WithError(err).Error("get schema version fail")
 				request.callback(nil, request.msg, fmt.Errorf("get schema version fail, err: %w", err))
 				return
@@ -596,7 +598,7 @@ func (p *partitionProducer) internalSend(request *sendRequest) {
 
 	// if msg is too large and chunking is disabled
 	if checkSize > maxMessageSize && !p.options.EnableChunking {
-		p.publishSemaphore.Release()
+		p.releaseSemaphoreAndMem(uncompressedPayloadSize)
 		request.callback(nil, request.msg, errMessageTooLarge)
 		p.log.WithError(errMessageTooLarge).
 			WithField("size", checkSize).
@@ -615,7 +617,7 @@ func (p *partitionProducer) internalSend(request *sendRequest) {
 	} else {
 		payloadChunkSize = int(p._getConn().GetMaxMessageSize()) - proto.Size(mm)
 		if payloadChunkSize <= 0 {
-			p.publishSemaphore.Release()
+			p.releaseSemaphoreAndMem(uncompressedPayloadSize)
 			request.callback(nil, msg, errMetaTooLarge)
 			p.log.WithError(errMetaTooLarge).
 				WithField("metadata size", proto.Size(mm)).
@@ -663,7 +665,8 @@ func (p *partitionProducer) internalSend(request *sendRequest) {
 					chunkRecorder:    cr,
 				}
 				// the permit of first chunk has acquired
-				if chunkID != 0 && !p.canAddToQueue(nsr) {
+				if chunkID != 0 && !p.canAddToQueue(nsr, 0) {
+					p.releaseSemaphoreAndMem(uncompressedPayloadSize - int64(rhs))
 					return
 				}
 				p.internalSingleSend(mm, compressedPayload[lhs:rhs], nsr, uint32(maxMessageSize))
@@ -688,7 +691,7 @@ func (p *partitionProducer) internalSend(request *sendRequest) {
 			// after flushing try again to add the current payload
 			if ok := p.batchBuilder.Add(smm, p.sequenceIDGenerator, uncompressedPayload, request,
 				msg.ReplicationClusters, deliverAt, schemaVersion, multiSchemaEnabled); !ok {
-				p.publishSemaphore.Release()
+				p.releaseSemaphoreAndMem(uncompressedPayloadSize)
 				request.callback(nil, request.msg, errFailAddToBatch)
 				p.log.WithField("size", uncompressedSize).
 					WithField("properties", msg.Properties).
@@ -797,7 +800,7 @@ func (p *partitionProducer) internalSingleSend(mm *pb.MessageMetadata,
 		maxMessageSize,
 	); err != nil {
 		request.callback(nil, request.msg, err)
-		p.publishSemaphore.Release()
+		p.releaseSemaphoreAndMem(int64(len(msg.Payload)))
 		p.log.WithError(err).Errorf("Single message serialize failed %s", msg.Value)
 		return
 	}
@@ -935,7 +938,7 @@ func (p *partitionProducer) failTimeoutMessages() {
 				sr := i.(*sendRequest)
 				if sr.msg != nil {
 					size := len(sr.msg.Payload)
-					p.publishSemaphore.Release()
+					p.releaseSemaphoreAndMem(int64(size))
 					p.metrics.MessagesPending.Dec()
 					p.metrics.BytesPending.Sub(float64(size))
 					p.metrics.PublishErrorsTimeout.Inc()
@@ -1139,8 +1142,7 @@ func (p *partitionProducer) ReceivedSendReceipt(response *pb.CommandSendReceipt)
 			sr := i.(*sendRequest)
 			if sr.msg != nil {
 				atomic.StoreInt64(&p.lastSequenceID, int64(pi.sequenceID))
-				p.publishSemaphore.Release()
-
+				p.releaseSemaphoreAndMem(int64(len(sr.msg.Payload)))
 				p.metrics.PublishLatency.Observe(float64(now-sr.publishTime.UnixNano()) / 1.0e9)
 				p.metrics.MessagesPublished.Inc()
 				p.metrics.MessagesPending.Dec()
@@ -1326,7 +1328,12 @@ func (p *partitionProducer) _getConn() internal.Connection {
 	return p.conn.Load().(internal.Connection)
 }
 
-func (p *partitionProducer) canAddToQueue(sr *sendRequest) bool {
+func (p *partitionProducer) releaseSemaphoreAndMem(size int64) {
+	p.publishSemaphore.Release()
+	p.client.memLimit.ReleaseMemory(size)
+}
+
+func (p *partitionProducer) canAddToQueue(sr *sendRequest, uncompressedPayloadSize int64) bool {
 	if p.options.DisableBlockIfQueueFull {
 		if !p.publishSemaphore.TryAcquire() {
 			if sr.callback != nil {
@@ -1334,9 +1341,24 @@ func (p *partitionProducer) canAddToQueue(sr *sendRequest) bool {
 			}
 			return false
 		}
-	} else if !p.publishSemaphore.Acquire(sr.ctx) {
-		sr.callback(nil, sr.msg, errContextExpired)
-		return false
+		if !p.client.memLimit.TryReserveMemory(uncompressedPayloadSize) {
+			p.publishSemaphore.Release()
+			if sr.callback != nil {
+				sr.callback(nil, sr.msg, errMemoryBufferIsFull)
+			}
+			return false
+		}
+
+	} else {
+		if !p.publishSemaphore.Acquire(sr.ctx) {
+			sr.callback(nil, sr.msg, errContextExpired)
+			return false
+		}
+		if !p.client.memLimit.ReserveMemory(sr.ctx, uncompressedPayloadSize) {
+			p.publishSemaphore.Release()
+			sr.callback(nil, sr.msg, errContextExpired)
+			return false
+		}
 	}
 	p.metrics.MessagesPending.Inc()
 	p.metrics.BytesPending.Add(float64(len(sr.msg.Payload)))


Test patch:
diff --git a/pulsar/internal/channel_cond_test.go b/pulsar/internal/channel_cond_test.go
new file mode 100644
index 0000000..a73d44e
--- /dev/null
+++ b/pulsar/internal/channel_cond_test.go
@@ -0,0 +1,55 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package internal
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+)
+
+func TestChCond(t *testing.T) {
+	cond := newCond(&sync.Mutex{})
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		cond.L.Lock()
+		cond.wait()
+		cond.L.Unlock()
+		wg.Done()
+	}()
+	time.Sleep(10 * time.Millisecond)
+	cond.broadcast()
+	wg.Wait()
+}
+
+func TestChCondWithContext(t *testing.T) {
+	cond := newCond(&sync.Mutex{})
+	wg := sync.WaitGroup{}
+	ctx, cancel := context.WithCancel(context.Background())
+	wg.Add(1)
+	go func() {
+		cond.L.Lock()
+		cond.waitWithContext(ctx)
+		cond.L.Unlock()
+		wg.Done()
+	}()
+	cancel()
+	wg.Wait()
+}
diff --git a/pulsar/internal/memory_limit_controller_test.go b/pulsar/internal/memory_limit_controller_test.go
new file mode 100644
index 0000000..a62c6e6
--- /dev/null
+++ b/pulsar/internal/memory_limit_controller_test.go
@@ -0,0 +1,184 @@
+// Licensed to the Apache Software Foundation (ASF) under one
+// or more contributor license agreements.  See the NOTICE file
+// distributed with this work for additional information
+// regarding copyright ownership.  The ASF licenses this file
+// to you under the Apache License, Version 2.0 (the
+// "License"); you may not use this file except in compliance
+// with the License.  You may obtain a copy of the License at
+//
+//   http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing,
+// software distributed under the License is distributed on an
+// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+// KIND, either express or implied.  See the License for the
+// specific language governing permissions and limitations
+// under the License.
+
+package internal
+
+import (
+	"context"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestLimit(t *testing.T) {
+
+	mlc := NewMemoryLimitController(100)
+
+	for i := 0; i < 101; i++ {
+		assert.True(t, mlc.TryReserveMemory(1))
+	}
+
+	assert.False(t, mlc.TryReserveMemory(1))
+	assert.Equal(t, int64(101), mlc.CurrentUsage())
+	assert.InDelta(t, 1.01, mlc.CurrentUsagePercent(), 0.000001)
+
+	mlc.ReleaseMemory(1)
+	assert.Equal(t, int64(100), mlc.CurrentUsage())
+	assert.InDelta(t, 1.0, mlc.CurrentUsagePercent(), 0.000001)
+
+	assert.True(t, mlc.TryReserveMemory(1))
+	assert.Equal(t, int64(101), mlc.CurrentUsage())
+
+	mlc.ForceReserveMemory(99)
+	assert.False(t, mlc.TryReserveMemory(1))
+	assert.Equal(t, int64(200), mlc.CurrentUsage())
+	assert.InDelta(t, 2.0, mlc.CurrentUsagePercent(), 0.000001)
+
+	mlc.ReleaseMemory(50)
+	assert.False(t, mlc.TryReserveMemory(1))
+	assert.Equal(t, int64(150), mlc.CurrentUsage())
+	assert.InDelta(t, 1.5, mlc.CurrentUsagePercent(), 0.000001)
+}
+
+func TestDisableLimit(t *testing.T) {
+	mlc := NewMemoryLimitController(-1)
+	assert.True(t, mlc.TryReserveMemory(1000000))
+	assert.True(t, mlc.ReserveMemory(context.Background(), 1000000))
+	mlc.ReleaseMemory(1000000)
+	assert.Equal(t, int64(1000000), mlc.CurrentUsage())
+}
+
+func TestMultiGoroutineTryReserveMem(t *testing.T) {
+	mlc := NewMemoryLimitController(10000)
+
+	// Multi goroutine try reserve memory.
+	wg := sync.WaitGroup{}
+
+	wg.Add(10)
+	for i := 0; i < 10; i++ {
+		go func() {
+			for i := 0; i < 1000; i++ {
+				assert.True(t, mlc.TryReserveMemory(1))
+			}
+			wg.Done()
+		}()
+	}
+	assert.True(t, mlc.TryReserveMemory(1))
+	wg.Wait()
+	assert.False(t, mlc.TryReserveMemory(1))
+	assert.Equal(t, int64(10001), mlc.CurrentUsage())
+	assert.InDelta(t, 1.0001, mlc.CurrentUsagePercent(), 0.000001)
+}
+
+func TestReserveWithContext(t *testing.T) {
+	mlc := NewMemoryLimitController(100)
+	assert.True(t, mlc.TryReserveMemory(101))
+	gorNum := 10
+
+	// Reserve ctx timeout
+	waitGroup := sync.WaitGroup{}
+	waitGroup.Add(gorNum)
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	defer cancel()
+	for i := 0; i < gorNum; i++ {
+		go func() {
+			assert.False(t, mlc.ReserveMemory(ctx, 1))
+			waitGroup.Done()
+		}()
+	}
+	waitGroup.Wait()
+	assert.Equal(t, int64(101), mlc.CurrentUsage())
+
+	// Reserve ctx cancel
+	waitGroup.Add(gorNum)
+	cancelCtx, cancel := context.WithCancel(context.Background())
+	for i := 0; i < gorNum; i++ {
+		go func() {
+			assert.False(t, mlc.ReserveMemory(cancelCtx, 1))
+			waitGroup.Done()
+		}()
+	}
+	cancel()
+	waitGroup.Wait()
+	assert.Equal(t, int64(101), mlc.CurrentUsage())
+}
+
+func TestBlocking(t *testing.T) {
+	mlc := NewMemoryLimitController(100)
+	assert.True(t, mlc.TryReserveMemory(101))
+	assert.Equal(t, int64(101), mlc.CurrentUsage())
+	assert.InDelta(t, 1.01, mlc.CurrentUsagePercent(), 0.000001)
+
+	gorNum := 10
+	chs := make([]chan int, gorNum)
+	for i := 0; i < gorNum; i++ {
+		chs[i] = make(chan int, 1)
+		go reserveMemory(mlc, chs[i])
+	}
+
+	// The threads are blocked since the quota is full
+	for i := 0; i < gorNum; i++ {
+		assert.False(t, awaitCh(chs[i]))
+	}
+	assert.Equal(t, int64(101), mlc.CurrentUsage())
+
+	mlc.ReleaseMemory(int64(gorNum))
+	for i := 0; i < gorNum; i++ {
+		assert.True(t, awaitCh(chs[i]))
+	}
+	assert.Equal(t, int64(101), mlc.CurrentUsage())
+}
+
+func TestStepRelease(t *testing.T) {
+	mlc := NewMemoryLimitController(100)
+	assert.True(t, mlc.TryReserveMemory(101))
+	assert.Equal(t, int64(101), mlc.CurrentUsage())
+	assert.InDelta(t, 1.01, mlc.CurrentUsagePercent(), 0.000001)
+
+	gorNum := 10
+	ch := make(chan int, 1)
+	for i := 0; i < gorNum; i++ {
+		go reserveMemory(mlc, ch)
+	}
+
+	// The threads are blocked since the quota is full
+	assert.False(t, awaitCh(ch))
+	assert.Equal(t, int64(101), mlc.CurrentUsage())
+
+	for i := 0; i < gorNum; i++ {
+		mlc.ReleaseMemory(1)
+		assert.True(t, awaitCh(ch))
+		assert.False(t, awaitCh(ch))
+	}
+	assert.Equal(t, int64(101), mlc.CurrentUsage())
+}
+
+func reserveMemory(mlc MemoryLimitController, ch chan int) {
+	mlc.ReserveMemory(context.Background(), 1)
+	ch <- 1
+}
+
+func awaitCh(ch chan int) bool {
+	select {
+	case <-ch:
+		return true
+	case <-time.After(100 * time.Millisecond):
+		return false
+	}
+}
diff --git a/pulsar/producer_test.go b/pulsar/producer_test.go
index e69a14c..f86d01a 100644
--- a/pulsar/producer_test.go
+++ b/pulsar/producer_test.go
@@ -1779,3 +1779,129 @@ func TestWaitForExclusiveProducer(t *testing.T) {
 	producer1.Close()
 	wg.Wait()
 }
+
+func TestMemLimitRejectProducerMessages(t *testing.T) {
+
+	c, err := NewClient(ClientOptions{
+		URL:              serviceURL,
+		MemoryLimitBytes: 100 * 1024,
+	})
+	assert.NoError(t, err)
+	defer c.Close()
+
+	topicName := newTopicName()
+	producer1, _ := c.CreateProducer(ProducerOptions{
+		Topic:                   topicName,
+		DisableBlockIfQueueFull: true,
+		DisableBatching:         false,
+		BatchingMaxPublishDelay: 100 * time.Second,
+		SendTimeout:             2 * time.Second,
+	})
+
+	producer2, _ := c.CreateProducer(ProducerOptions{
+		Topic:                   topicName,
+		DisableBlockIfQueueFull: true,
+		DisableBatching:         false,
+		BatchingMaxPublishDelay: 100 * time.Second,
+		SendTimeout:             2 * time.Second,
+	})
+
+	n := 101
+	for i := 0; i < n/2; i++ {
+		producer1.SendAsync(context.Background(), &ProducerMessage{
+			Payload: make([]byte, 1024),
+		}, func(id MessageID, message *ProducerMessage, e error) {})
+
+		producer2.SendAsync(context.Background(), &ProducerMessage{
+			Payload: make([]byte, 1024),
+		}, func(id MessageID, message *ProducerMessage, e error) {})
+	}
+	// Last message in order to reach the limit
+	producer1.SendAsync(context.Background(), &ProducerMessage{
+		Payload: make([]byte, 1024),
+	}, func(id MessageID, message *ProducerMessage, e error) {})
+	time.Sleep(100 * time.Millisecond)
+	assert.Equal(t, int64(n*1024), c.(*client).memLimit.CurrentUsage())
+
+	_, err = producer1.Send(context.Background(), &ProducerMessage{
+		Payload: make([]byte, 1024),
+	})
+	assert.Error(t, err)
+	assert.ErrorContains(t, err, getResultStr(ClientMemoryBufferIsFull))
+
+	_, err = producer2.Send(context.Background(), &ProducerMessage{
+		Payload: make([]byte, 1024),
+	})
+	assert.Error(t, err)
+	assert.ErrorContains(t, err, getResultStr(ClientMemoryBufferIsFull))
+
+	// flush pending msg
+	err = producer1.Flush()
+	assert.NoError(t, err)
+	err = producer2.Flush()
+	assert.NoError(t, err)
+	assert.Equal(t, int64(0), c.(*client).memLimit.CurrentUsage())
+
+	_, err = producer1.Send(context.Background(), &ProducerMessage{
+		Payload: make([]byte, 1024),
+	})
+	assert.NoError(t, err)
+	_, err = producer2.Send(context.Background(), &ProducerMessage{
+		Payload: make([]byte, 1024),
+	})
+	assert.NoError(t, err)
+}
+
+func TestMemLimitContextCancel(t *testing.T) {
+
+	c, err := NewClient(ClientOptions{
+		URL:              serviceURL,
+		MemoryLimitBytes: 100 * 1024,
+	})
+	assert.NoError(t, err)
+	defer c.Close()
+
+	topicName := newTopicName()
+	producer, _ := c.CreateProducer(ProducerOptions{
+		Topic:                   topicName,
+		DisableBlockIfQueueFull: false,
+		DisableBatching:         false,
+		BatchingMaxPublishDelay: 100 * time.Second,
+		SendTimeout:             2 * time.Second,
+	})
+
+	n := 101
+	ctx, cancel := context.WithCancel(context.Background())
+	for i := 0; i < n; i++ {
+		producer.SendAsync(ctx, &ProducerMessage{
+			Payload: make([]byte, 1024),
+		}, func(id MessageID, message *ProducerMessage, e error) {})
+	}
+	time.Sleep(100 * time.Millisecond)
+	assert.Equal(t, int64(n*1024), c.(*client).memLimit.CurrentUsage())
+
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+	go func() {
+		producer.SendAsync(ctx, &ProducerMessage{
+			Payload: make([]byte, 1024),
+		}, func(id MessageID, message *ProducerMessage, e error) {
+			assert.Error(t, e)
+			assert.ErrorContains(t, e, getResultStr(TimeoutError))
+			wg.Done()
+		})
+	}()
+
+	// cancel pending msg
+	cancel()
+	wg.Wait()
+
+	err = producer.Flush()
+	assert.NoError(t, err)
+	assert.Equal(t, int64(0), c.(*client).memLimit.CurrentUsage())
+
+	_, err = producer.Send(context.Background(), &ProducerMessage{
+		Payload: make([]byte, 1024),
+	})
+	assert.NoError(t, err)
+}