Provide lock-free access to partition consumer compression providers (#689)

The compression providers map in the partition consumer is a textbook case
for using go's lock-free sync.Map: the set of map entries is stable and
access is read-only.

On machines with 4 cores or greater, read contention on the sync.RWMutex
outweighs the cost of using a sync.Map.

Below is an old article on the subject, but it still holds true today:
https://medium.com/@deckarep/the-new-kid-in-town-gos-sync-map-de24a6bf7c2c

Signed-off-by: Daniel Ferstay <dferstay@splunk.com>

Co-authored-by: Daniel Ferstay <dferstay@splunk.com>

Author: dferstay

pulsar/consumer_partition.go

@@ -144,8 +144,7 @@ type partitionConsumer struct {
 
 	log log.Logger
 
-	providersMutex       sync.RWMutex
-	compressionProviders map[pb.CompressionType]compression.Provider
+	compressionProviders sync.Map //map[pb.CompressionType]compression.Provider
 	metrics              *internal.LeveledMetrics
 	decryptor            cryptointernal.Decryptor
 }
@@ -171,7 +170,7 @@ func newPartitionConsumer(parent Consumer, client *client, options *partitionCon
 		closeCh:              make(chan struct{}),
 		clearQueueCh:         make(chan func(id trackingMessageID)),
 		clearMessageQueuesCh: make(chan chan struct{}),
-		compressionProviders: make(map[pb.CompressionType]compression.Provider),
+		compressionProviders: sync.Map{},
 		dlq:                  dlq,
 		metrics:              metrics,
 	}
@@ -967,11 +966,15 @@ func (pc *partitionConsumer) internalClose(req *closeRequest) {
 		pc.log.Info("Closed consumer")
 	}
 
-	pc.providersMutex.Lock()
-	for _, provider := range pc.compressionProviders {
-		provider.Close()
-	}
-	pc.providersMutex.Unlock()
+	pc.compressionProviders.Range(func(_, v interface{}) bool {
+		if provider, ok := v.(compression.Provider); ok {
+			provider.Close()
+		} else {
+			err := fmt.Errorf("unexpected compression provider type: %T", v)
+			pc.log.WithError(err).Warn("Failed to close compression provider")
+		}
+		return true
+	})
 
 	pc.setConsumerState(consumerClosed)
 	pc._getConn().DeleteConsumeHandler(pc.consumerID)
@@ -1192,19 +1195,26 @@ func getPreviousMessage(mid trackingMessageID) trackingMessageID {
 }
 
 func (pc *partitionConsumer) Decompress(msgMeta *pb.MessageMetadata, payload internal.Buffer) (internal.Buffer, error) {
-	pc.providersMutex.RLock()
-	provider, ok := pc.compressionProviders[msgMeta.GetCompression()]
-	pc.providersMutex.RUnlock()
+	providerEntry, ok := pc.compressionProviders.Load(msgMeta.GetCompression())
 	if !ok {
-		var err error
-		if provider, err = pc.initializeCompressionProvider(msgMeta.GetCompression()); err != nil {
+		newProvider, err := pc.initializeCompressionProvider(msgMeta.GetCompression())
+		if err != nil {
 			pc.log.WithError(err).Error("Failed to decompress message.")
 			return nil, err
 		}
 
-		pc.providersMutex.Lock()
-		pc.compressionProviders[msgMeta.GetCompression()] = provider
-		pc.providersMutex.Unlock()
+		var loaded bool
+		providerEntry, loaded = pc.compressionProviders.LoadOrStore(msgMeta.GetCompression(), newProvider)
+		if loaded {
+			// another thread already loaded this provider, so close the one we just initialized
+			newProvider.Close()
+		}
+	}
+	provider, ok := providerEntry.(compression.Provider)
+	if !ok {
+		err := fmt.Errorf("unexpected compression provider type: %T", providerEntry)
+		pc.log.WithError(err).Error("Failed to decompress message.")
+		return nil, err
 	}
 
 	uncompressed, err := provider.Decompress(nil, payload.ReadableSlice(), int(msgMeta.GetUncompressedSize()))

pulsar/consumer_partition_test.go

@@ -18,11 +18,10 @@
 package pulsar
 
 import (
+	"sync"
 	"testing"
 
-	"github.com/apache/pulsar-client-go/pulsar/internal/compression"
 	"github.com/apache/pulsar-client-go/pulsar/internal/crypto"
-	pb "github.com/apache/pulsar-client-go/pulsar/internal/pulsar_proto"
 
 	"github.com/stretchr/testify/assert"
 
@@ -34,7 +33,7 @@ func TestSingleMessageIDNoAckTracker(t *testing.T) {
 	pc := partitionConsumer{
 		queueCh:              make(chan []*message, 1),
 		eventsCh:             eventsCh,
-		compressionProviders: make(map[pb.CompressionType]compression.Provider),
+		compressionProviders: sync.Map{},
 		options:              &partitionConsumerOpts{},
 		metrics:              internal.NewMetricsProvider(4, map[string]string{}).GetLeveledMetrics("topic"),
 		decryptor:            crypto.NewNoopDecryptor(),
@@ -66,7 +65,7 @@ func TestBatchMessageIDNoAckTracker(t *testing.T) {
 	pc := partitionConsumer{
 		queueCh:              make(chan []*message, 1),
 		eventsCh:             eventsCh,
-		compressionProviders: make(map[pb.CompressionType]compression.Provider),
+		compressionProviders: sync.Map{},
 		options:              &partitionConsumerOpts{},
 		metrics:              internal.NewMetricsProvider(4, map[string]string{}).GetLeveledMetrics("topic"),
 		decryptor:            crypto.NewNoopDecryptor(),
@@ -98,7 +97,7 @@ func TestBatchMessageIDWithAckTracker(t *testing.T) {
 	pc := partitionConsumer{
 		queueCh:              make(chan []*message, 1),
 		eventsCh:             eventsCh,
-		compressionProviders: make(map[pb.CompressionType]compression.Provider),
+		compressionProviders: sync.Map{},
 		options:              &partitionConsumerOpts{},
 		metrics:              internal.NewMetricsProvider(4, map[string]string{}).GetLeveledMetrics("topic"),
 		decryptor:            crypto.NewNoopDecryptor(),