Support snappy compression (#1406)

Author: geniusjoe

pulsar/consumer_partition.go

@@ -2242,6 +2242,8 @@ func (pc *partitionConsumer) initializeCompressionProvider(
 		return compression.NewLz4Provider(), nil
 	case pb.CompressionType_ZSTD:
 		return compression.NewZStdProvider(compression.Default), nil
+	case pb.CompressionType_SNAPPY:
+		return compression.NewSnappyProvider(), nil
 	}
 
 	return nil, fmt.Errorf("unsupported compression type: %v", compressionType)

pulsar/consumer_test.go

@@ -1293,6 +1293,98 @@ func TestConsumerCompression(t *testing.T) {
 	}
 }
 
+func TestConsumerMultiCompressions(t *testing.T) {
+	type testProvider struct {
+		name            string
+		compressionType CompressionType
+	}
+
+	providers := []testProvider{
+		{"zlib", ZLib},
+		{"lz4", LZ4},
+		{"zstd", ZSTD},
+		{"snappy", SNAPPY},
+	}
+
+	for _, provider := range providers {
+		p := provider
+		t.Run(p.name, func(t *testing.T) {
+			client, err := NewClient(ClientOptions{
+				URL: lookupURL,
+			})
+
+			assert.Nil(t, err)
+			defer client.Close()
+
+			batchTopic, nonBatchTopic := newTopicName(), newTopicName()
+			ctx := context.Background()
+
+			// enable batching
+			batchProducer, err := client.CreateProducer(ProducerOptions{
+				Topic:           batchTopic,
+				CompressionType: p.compressionType,
+				DisableBatching: false,
+			})
+			assert.Nil(t, err)
+			defer batchProducer.Close()
+
+			batchConsumer, err := client.Subscribe(ConsumerOptions{
+				Topic:            batchTopic,
+				SubscriptionName: "sub-1",
+			})
+			assert.Nil(t, err)
+			defer batchConsumer.Close()
+
+			const N = 100
+			for i := 0; i < N; i++ {
+				batchProducer.SendAsync(ctx, &ProducerMessage{
+					Payload: []byte(fmt.Sprintf("msg-content-%d-batching-enabled", i)),
+				}, func(_ MessageID, _ *ProducerMessage, err error) {
+					assert.Nil(t, err)
+				})
+			}
+
+			for i := 0; i < N; i++ {
+				msg, err := batchConsumer.Receive(ctx)
+				assert.Nil(t, err)
+				assert.Equal(t, fmt.Sprintf("msg-content-%d-batching-enabled", i), string(msg.Payload()))
+				batchConsumer.Ack(msg)
+			}
+
+			// disable batching
+			nonBatchProducer, err := client.CreateProducer(ProducerOptions{
+				Topic:           nonBatchTopic,
+				CompressionType: p.compressionType,
+				DisableBatching: true,
+			})
+			assert.Nil(t, err)
+			defer nonBatchProducer.Close()
+
+			nonBatchConsumer, err := client.Subscribe(ConsumerOptions{
+				Topic:            nonBatchTopic,
+				SubscriptionName: "sub-1",
+			})
+			assert.Nil(t, err)
+			defer nonBatchConsumer.Close()
+
+			for i := 0; i < N; i++ {
+				if _, err := nonBatchProducer.Send(ctx, &ProducerMessage{
+					Payload: []byte(fmt.Sprintf("msg-content-%d-batching-disabled", i)),
+				}); err != nil {
+					t.Fatal(err)
+				}
+			}
+
+			for i := 0; i < N; i++ {
+				msg, err := nonBatchConsumer.Receive(ctx)
+				assert.Nil(t, err)
+				assert.Equal(t, fmt.Sprintf("msg-content-%d-batching-disabled", i), string(msg.Payload()))
+				nonBatchConsumer.Ack(msg)
+			}
+		})
+	}
+}
+
 func TestConsumerCompressionWithBatches(t *testing.T) {
 	client, err := NewClient(ClientOptions{
 		URL: lookupURL,

pulsar/internal/batch_builder.go

@@ -313,6 +313,8 @@ func GetCompressionProvider(
 		return compression.NewZLibProvider()
 	case pb.CompressionType_ZSTD:
 		return compression.NewZStdProvider(level)
+	case pb.CompressionType_SNAPPY:
+		return compression.NewSnappyProvider()
 	default:
 		panic("unsupported compression type")
 	}

pulsar/internal/compression/compression_bench_test.go

@@ -72,6 +72,7 @@ var benchmarkProviders = []testProvider{
 	{"zstd-cgo-level-fastest", newCGoZStdProvider(Faster), nil},
 	{"zstd-cgo-level-default", newCGoZStdProvider(Default), nil},
 	{"zstd-cgo-level-best", newCGoZStdProvider(Better), nil},
+	{"snappy", NewSnappyProvider(), nil},
 }
 
 func BenchmarkCompression(b *testing.B) {

pulsar/internal/compression/compression_test.go

@@ -36,6 +36,7 @@ var providers = []testProvider{
 	{"lz4", NewLz4Provider(), []byte{0x50, 0x68, 0x65, 0x6c, 0x6c, 0x6f}},
 	{"zstd", NewZStdProvider(Default),
 		[]byte{0x28, 0xb5, 0x2f, 0xfd, 0x20, 0x05, 0x29, 0x00, 0x00, 0x68, 0x65, 0x6c, 0x6c, 0x6f}},
+	{"snappy", NewSnappyProvider(), []byte{0x05, 0x10, 0x68, 0x65, 0x6c, 0x6c, 0x6f}},
 }
 
 func TestCompression(t *testing.T) {

pulsar/internal/compression/snappy.go

@@ -0,0 +1,50 @@
+// 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 compression
+
+import (
+	"github.com/klauspost/compress/snappy"
+)
+
+type snappyProvider struct {
+}
+
+// NewSnappyProvider returns a Provider interface
+func NewSnappyProvider() Provider {
+	return &snappyProvider{}
+}
+
+func (p *snappyProvider) CompressMaxSize(srcSize int) int {
+	return snappy.MaxEncodedLen(srcSize)
+}
+
+func (p *snappyProvider) Compress(dst, src []byte) []byte {
+	return snappy.Encode(dst, src)
+}
+
+func (p *snappyProvider) Decompress(dst, src []byte, _ int) ([]byte, error) {
+	return snappy.Decode(dst, src)
+}
+
+func (p *snappyProvider) Close() error {
+	return nil
+}
+
+func (p *snappyProvider) Clone() Provider {
+	return &snappyProvider{}
+}

pulsar/producer.go

@@ -40,6 +40,7 @@ const (
 	LZ4
 	ZLib
 	ZSTD
+	SNAPPY
 )
 
 type CompressionLevel int
@@ -120,6 +121,7 @@ type ProducerOptions struct {
 	//  - LZ4
 	//  - ZLIB
 	//  - ZSTD
+	//	- SNAPPY
 	//
 	// Note: ZSTD is supported since Pulsar 2.3. Consumers will need to be at least at that
 	// release in order to be able to receive messages compressed with ZSTD.

pulsar/producer_test.go

@@ -245,6 +245,7 @@ func TestProducerCompression(t *testing.T) {
 		{"zlib", ZLib},
 		{"lz4", LZ4},
 		{"zstd", ZSTD},
+		{"snappy", SNAPPY},
 	}
 
 	for _, provider := range providers {