fix race condition in lz4_raw compression

xitongsys · Aug 17, 2024 · 91a1c46 · 91a1c46
1 parent 9ddb178
commit 91a1c46
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 6 deletions.
diff --git a/compress/lz4_raw.go b/compress/lz4_raw.go
@@ -9,11 +9,11 @@ import (
 )
 
 func init() {
-	lz4hc := lz4.CompressorHC{
-		Level: lz4.CompressionLevel(9),
-	}
 	compressors[parquet.CompressionCodec_LZ4_RAW] = &Compressor{
 		Compress: func(buf []byte) []byte {
+			lz4hc := lz4.CompressorHC{
+				Level: lz4.CompressionLevel(9),
+			}
 			res := make([]byte, lz4.CompressBlockBound(len(buf)))
 			count, _ := lz4hc.CompressBlock(buf, res)
 			return res[:count]

diff --git a/compress/lz4_raw_test.go b/compress/lz4_raw_test.go
@@ -2,6 +2,7 @@ package compress
 
 import (
 	"bytes"
+	"sync"
 	"testing"
 
 	"github.com/xitongsys/parquet-go/parquet"
@@ -15,10 +16,18 @@ func TestLz4RawCompress(t *testing.T) {
 	}
 
 	// compression
-	output := lz4RawCompressor.Compress(input)
-	if !bytes.Equal(compressed, output) {
-		t.Fatalf("expected output %s but was %s", string(compressed), string(output))
+	var wg sync.WaitGroup
+	for i := 0; i < 10; i++ {
+		go func() {
+			wg.Add(1)
+			defer wg.Done()
+			output := lz4RawCompressor.Compress(input)
+			if !bytes.Equal(compressed, output) {
+				t.Fatalf("expected output %s but was %s", string(compressed), string(output))
+			}
+		}()
 	}
+	wg.Wait()
 
 	// uncompression
 	output, err := lz4RawCompressor.Uncompress(compressed)