Add multi-threading support for parallel compression

Makefile

@@ -3,58 +3,78 @@ GOARCH ?= $(shell go env GOARCH)
 GOOS_GOARCH := $(GOOS)_$(GOARCH)
 GOOS_GOARCH_NATIVE := $(shell go env GOHOSTOS)_$(shell go env GOHOSTARCH)
 LIBZSTD_NAME := libzstd_$(GOOS_GOARCH).a
-ZSTD_VERSION ?= v1.5.7
-ZIG_BUILDER_IMAGE=euantorano/zig:0.10.1
-BUILDER_IMAGE := local/builder_musl:2.0.0-$(shell echo $(ZIG_BUILDER_IMAGE) | tr : _ | tr / _)-1
+ZSTD_VERSION ?= v1.5.7-kernel
+ZIG_BUILDER_IMAGE := euantorano/zig:0.10.1
 
-.PHONY: libzstd.a $(LIBZSTD_NAME)
+# Detect available container runtime
+CONTAINER_RUNTIME := $(shell \
+	if command -v docker >/dev/null 2>&1; then \
+		echo "docker"; \
+	elif command -v nerdctl >/dev/null 2>&1; then \
+		echo "nerdctl"; \
+	elif command -v podman >/dev/null 2>&1; then \
+		echo "podman"; \
+	else \
+		echo "none"; \
+	fi)
+
+# Show which runtime is being used
+$(info Using container runtime: $(CONTAINER_RUNTIME))
+
+# Parallel compilation flags
+JOBS := $(shell nproc 2>/dev/null || sysctl -n hw.ncpu 2>/dev/null || echo 4)
+MAKEFLAGS += -j$(JOBS)
+
+.PHONY: libzstd.a $(LIBZSTD_NAME) test test-no-lto bench bench-no-lto
 
 libzstd.a: $(LIBZSTD_NAME)
 $(LIBZSTD_NAME):
 ifeq ($(GOOS_GOARCH),$(GOOS_GOARCH_NATIVE))
 	rm -f $(LIBZSTD_NAME)
-	cd zstd/lib && ZSTD_LEGACY_SUPPORT=0 MOREFLAGS=$(MOREFLAGS) $(MAKE) clean libzstd.a
+	cd zstd/lib && ZSTD_LEGACY_SUPPORT=0 AR="gcc-ar" ARFLAGS="rcs" MOREFLAGS="-DZSTD_MULTITHREAD=1 -O3 -flto=1 $(MOREFLAGS)" LDFLAGS="-flto=1 -fuse-linker-plugin" $(MAKE) clean libzstd.a
 	mv zstd/lib/libzstd.a $(LIBZSTD_NAME)
 else ifeq ($(GOOS_GOARCH),linux_amd64)
-	TARGET=x86_64-linux GOARCH=amd64 GOOS=linux $(MAKE) package-arch
+	TARGET=x86_64-linux GOARCH=amd64 GOOS=linux ARCH_FLAGS="-mcpu=x86_64+sse4_2+avx2+bmi2" $(MAKE) package-arch
 else ifeq ($(GOOS_GOARCH),linux_arm)
-	TARGET=arm-linux-gnueabi GOARCH=arm GOOS=linux $(MAKE) package-arch
+	TARGET=arm-linux-gnueabi GOARCH=arm GOOS=linux ARCH_FLAGS="-mcpu=generic" $(MAKE) package-arch
 else ifeq ($(GOOS_GOARCH),linux_arm64)
-	TARGET=aarch64-linux GOARCH=arm64 GOOS=linux $(MAKE) package-arch
+	TARGET=aarch64-linux GOARCH=arm64 GOOS=linux ARCH_FLAGS="-mcpu=generic" $(MAKE) package-arch
 else ifeq ($(GOOS_GOARCH),linux_ppc64le)
-	TARGET=x86_64-linux GOARCH=ppc64le GOOS=linux $(MAKE) package-arch
+	TARGET=x86_64-linux GOARCH=ppc64le GOOS=linux ARCH_FLAGS="" $(MAKE) package-arch
 else ifeq ($(GOOS_GOARCH),linux_musl_amd64)
-	TARGET=x86_64-linux-musl GOARCH=amd64 GOOS=linux_musl $(MAKE) package-arch
+	TARGET=x86_64-linux-musl GOARCH=amd64 GOOS=linux_musl ARCH_FLAGS="-mcpu=x86_64+sse4_2+avx2+bmi2" $(MAKE) package-arch
 else ifeq ($(GOOS_GOARCH),linux_musl_arm64)
-	TARGET=aarch64-linux-musl GOARCH=arm64 GOOS=linux_musl $(MAKE) package-arch
+	TARGET=aarch64-linux-musl GOARCH=arm64 GOOS=linux_musl ARCH_FLAGS="-mcpu=generic" $(MAKE) package-arch
 else ifeq ($(GOOS_GOARCH),darwin_arm64)
-	TARGET=aarch64-macos GOARCH=arm64 GOOS=darwin $(MAKE) package-arch
+	TARGET=aarch64-macos GOARCH=arm64 GOOS=darwin ARCH_FLAGS="-mcpu=apple_m1" $(MAKE) package-arch
 else ifeq ($(GOOS_GOARCH),darwin_amd64)
-	TARGET=x86_64-macos GOARCH=amd64 GOOS=darwin $(MAKE) package-arch
+	TARGET=x86_64-macos GOARCH=amd64 GOOS=darwin ARCH_FLAGS="-mcpu=x86_64+sse4_2+avx2" $(MAKE) package-arch
 else ifeq ($(GOOS_GOARCH),windows_amd64)
-	TARGET=x86_64-windows GOARCH=amd64 GOOS=windows GOARCH=amd64 $(MAKE) package-arch
+	TARGET=x86_64-windows GOARCH=amd64 GOOS=windows GOARCH=amd64 ARCH_FLAGS="-mcpu=x86_64+sse4_2+avx2+bmi2" $(MAKE) package-arch
 endif
 
-package-builder:
-	(docker image ls --format '{{.Repository}}:{{.Tag}}' | grep -q '$(BUILDER_IMAGE)$$') \
-		|| docker build \
-			--build-arg builder_image=$(ZIG_BUILDER_IMAGE) \
-			--tag $(BUILDER_IMAGE) \
-			builder
-
-package-arch: package-builder
+package-arch:
+ifeq ($(CONTAINER_RUNTIME),none)
+	$(error No container runtime found. Please install docker, nerdctl, or podman for cross-compilation)
+endif
 	rm -f $(LIBZSTD_NAME)
-	docker run --rm \
+	$(CONTAINER_RUNTIME) run --rm \
+		--entrypoint /bin/sh \
 		--mount type=bind,src="$(shell pwd)",dst=/zstd \
-		-w /zstd \
+		-w /zstd/zstd/lib \
 		$(DOCKER_OPTS) \
-		$(BUILDER_IMAGE) \
-		-c 'cd zstd/lib && \
+		$(ZIG_BUILDER_IMAGE) \
+		-c 'apk add --no-cache make && \
+			if echo "$(TARGET)" | grep -q "macos\|darwin"; then \
+				LTO_FLAG=""; \
+			else \
+				LTO_FLAG="-flto"; \
+			fi; \
 			ZSTD_LEGACY_SUPPORT=0 AR="zig ar" \
-			CC="zig cc -target $(TARGET)" \
-			CXX="zig cc -target $(TARGET)" \
-			MOREFLAGS=$(MOREFLAGS) \
-			$(MAKE) clean libzstd.a'
+			CC="zig cc -target $(TARGET) -O3 $$LTO_FLAG $(ARCH_FLAGS)" \
+			CXX="zig cc -target $(TARGET) -O3 $$LTO_FLAG $(ARCH_FLAGS)" \
+			MOREFLAGS="-DZSTD_MULTITHREAD=1 -O3 $$LTO_FLAG $(ARCH_FLAGS) $(MOREFLAGS)" \
+			make -j$(shell nproc 2>/dev/null || echo 4) clean libzstd.a'
 	mv -f zstd/lib/libzstd.a $(LIBZSTD_NAME)
 
 # freebsd and illumos aren't supported by zig compiler atm.
@@ -85,7 +105,13 @@ update-zstd:
 	$(MAKE) release
 
 test:
+	CGO_LDFLAGS_ALLOW='-flto.*' CGO_ENABLED=1 GOEXPERIMENT=cgocheck2 go test -ldflags="-linkmode=external" -v
+
+test-no-lto:
 	CGO_ENABLED=1 GOEXPERIMENT=cgocheck2 go test -v
 
 bench:
+	CGO_LDFLAGS_ALLOW='-flto.*' CGO_ENABLED=1 go test -ldflags="-linkmode=external" -bench=.
+
+bench-no-lto:
 	CGO_ENABLED=1 go test -bench=.

build.log

@@ -0,0 +1,39 @@
+Using container runtime: nerdctl
+rm -f libzstd_linux_amd64.a
+cd zstd/lib && ZSTD_LEGACY_SUPPORT=0 MOREFLAGS="-DZSTD_MULTITHREAD=1 -O3 -flto=1 -fwhole-program " LDFLAGS="-flto=1 -fuse-linker-plugin" make clean libzstd.a
+make[1]: Entering directory '/var/home/grigory/containerd/gozstd/zstd/lib'
+Cleaning library completed
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/debug.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/entropy_common.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/error_private.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/fse_decompress.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/pool.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/xxhash.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/threading.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_common.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/fse_compress.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/hist.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/huf_compress.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_compress.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_compress_literals.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_compress_sequences.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_compress_superblock.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_double_fast.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_fast.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_lazy.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_ldm.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_opt.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_preSplit.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstdmt_compress.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/huf_decompress.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_ddict.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_decompress.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zstd_decompress_block.o
+AS obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/huf_decompress_amd64.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/cover.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/divsufsort.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/fastcover.o
+CC obj/conf_4ac72047274d3e00ad8a036de02efe2b/static/zdict.o
+compiling single-threaded static library 1.5.7
+make[1]: Leaving directory '/var/home/grigory/containerd/gozstd/zstd/lib'
+mv zstd/lib/libzstd.a libzstd_linux_amd64.a

go.mod

@@ -1,3 +1,3 @@
 module github.com/valyala/gozstd
 
-go 1.12
+go 1.12

writer.go

@@ -64,6 +64,7 @@ type Writer struct {
 	w                io.Writer
 	compressionLevel int
 	wlog             int
+	nbWorkers        int
 	cs               *C.ZSTD_CStream
 	cd               *CDict
 
@@ -143,6 +144,12 @@ type WriterParams struct {
 
 	// Dict is optional dictionary used for compression.
 	Dict *CDict
+
+	// NbWorkers controls the number of compression threads.
+	// 0 = single-threaded mode (default)
+	// >=1 = multi-threaded mode with specified number of workers
+	// Note: The library must be compiled with ZSTD_MULTITHREAD=1 for this to work.
+	NbWorkers int
 }
 
 // NewWriterParams returns new zstd writer writing compressed data to w
@@ -174,6 +181,7 @@ func NewWriterParams(w io.Writer, params *WriterParams) *Writer {
 		w:                w,
 		compressionLevel: params.CompressionLevel,
 		wlog:             params.WindowLog,
+		nbWorkers:        params.NbWorkers,
 		cs:               cs,
 		cd:               params.Dict,
 		inBuf:            inBuf,
@@ -194,6 +202,7 @@ func (zw *Writer) Reset(w io.Writer, cd *CDict, compressionLevel int) {
 	params := WriterParams{
 		CompressionLevel: compressionLevel,
 		WindowLog:        zw.wlog,
+		NbWorkers:        zw.nbWorkers,
 		Dict:             cd,
 	}
 	zw.ResetWriterParams(w, &params)
@@ -206,6 +215,9 @@ func (zw *Writer) ResetWriterParams(w io.Writer, params *WriterParams) {
 	zw.outBuf.size = cstreamOutBufSize
 	zw.outBuf.pos = 0
 
+	zw.compressionLevel = params.CompressionLevel
+	zw.wlog = params.WindowLog
+	zw.nbWorkers = params.NbWorkers
 	zw.cd = params.Dict
 	initCStream(zw.cs, *params)
 
@@ -230,6 +242,13 @@ func initCStream(cs *C.ZSTD_CStream, params WriterParams) {
 		C.ZSTD_cParameter(C.ZSTD_c_windowLog),
 		C.int(params.WindowLog))
 	ensureNoError("ZSTD_CCtx_setParameter", result)
+
+	// Set number of workers for multi-threaded compression
+	result = C.ZSTD_CCtx_setParameter_wrapper(
+		C.uintptr_t(uintptr(unsafe.Pointer(cs))),
+		C.ZSTD_cParameter(400), // ZSTD_c_nbWorkers
+		C.int(params.NbWorkers))
+	ensureNoError("ZSTD_CCtx_setParameter", result)
 }
 
 func freeCStream(v interface{}) {

writer_test.go

@@ -496,3 +496,54 @@ func TestWriterBig(t *testing.T) {
 		t.Fatalf("unequal writtenBB and readBB\nwrittenBB=\n%X\nreadBB=\n%X", writtenBB.Bytes(), readBB.Bytes())
 	}
 }
+
+func TestWriterMultiThreading(t *testing.T) {
+	// Test data - make it large enough to benefit from multi-threading
+	data := make([]byte, 10*1024*1024) // 10MB
+	for i := range data {
+		data[i] = byte(i % 256)
+	}
+
+	// Test different numbers of workers
+	workerCounts := []int{0, 1, 2, 4, 8}
+
+	for _, nbWorkers := range workerCounts {
+		t.Run(fmt.Sprintf("NbWorkers=%d", nbWorkers), func(t *testing.T) {
+			params := &WriterParams{
+				CompressionLevel: 3,
+				NbWorkers:        nbWorkers,
+			}
+
+			var bb bytes.Buffer
+			zw := NewWriterParams(&bb, params)
+
+			// Write the data
+			n, err := zw.Write(data)
+			if err != nil {
+				t.Fatalf("cannot write data with nbWorkers=%d: %s", nbWorkers, err)
+			}
+			if n != len(data) {
+				t.Fatalf("unexpected bytes written: got %d, want %d", n, len(data))
+			}
+
+			// Close the writer
+			if err := zw.Close(); err != nil {
+				t.Fatalf("cannot close writer with nbWorkers=%d: %s", nbWorkers, err)
+			}
+			zw.Release()
+
+			// Decompress and verify
+			decompressed, err := Decompress(nil, bb.Bytes())
+			if err != nil {
+				t.Fatalf("cannot decompress data with nbWorkers=%d: %s", nbWorkers, err)
+			}
+
+			if !bytes.Equal(decompressed, data) {
+				t.Fatalf("decompressed data mismatch with nbWorkers=%d", nbWorkers)
+			}
+
+			t.Logf("nbWorkers=%d: compressed %d bytes to %d bytes (%.2f%% ratio)",
+				nbWorkers, len(data), bb.Len(), float64(bb.Len())/float64(len(data))*100)
+		})
+	}
+}