<commit_before>package network

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"io/ioutil"
	"math/rand"
	"sync"
	"testing"
	"time"
)

func BenchmarkGzip(b *testing.B) {
	b.Run("standard gzip", BenchmarkStandardGzip)
	b.Run("best compression gzip", BenchmarkBestCompGzip)
	b.Run("best speed gzip", BenchmarkBestSpeedGzip)
}

type gzipWriter struct {
	writer *gzip.Writer
	buffer *bytes.Buffer
}

func (w gzipWriter) Write(msg []byte) []byte {
	_, err := w.writer.Write(msg)
	if err != nil {
		return msg
	}

	err = w.writer.Flush()
	if err != nil {
		return msg
	}

	err = w.writer.Close()
	if err != nil {
		return msg
	}

	return w.buffer.Bytes()
}

func (w gzipWriter) Reset() {
	w.buffer.Reset()
	w.writer.Reset(w.buffer)
}

var pool = sync.Pool{
	New: func() interface{} {
		var buffer bytes.Buffer
		gWriter := gzip.NewWriter(&buffer)
		return gzipWriter{
			writer: gWriter,
			buffer: &buffer,
		}
	},
}

func BenchmarkPooledVsNonPooledGzip(b *testing.B) {
	b.Run("PooledGzip", BenchmarkPooledGzip)
	b.Run("NonPooledGzip", BenchmarkStandardGzipCompression)
}

var cmpBytes []byte

func BenchmarkPooledGzip(b *testing.B) {
	for i := 0; i < b.N; i++ {
		s := RandomString(2048)
		sBytes := []byte(s)
		t1 := time.Now()
		w := pool.Get().(gzipWriter)
		cmpBytes = w.Write(sBytes)
		w.Reset()
		pool.Put(w)
		b.ReportMetric(float64(time.Since(t1).Nanoseconds()), "CompressTime")
		b.ReportMetric(float64(len(sBytes)-len(cmpBytes)), "BytesSaved")
	}
}

func BenchmarkStandardGzipCompression(b *testing.B) {
	for i := 0; i < b.N; i++ {
		s := RandomString(1024)
		sBytes := []byte(s)
		t1 := time.Now()
		cmpBytes = standardCompress(sBytes, gzip.DefaultCompression)
		b.ReportMetric(float64(time.Since(t1).Nanoseconds()), "CompressTime")
		b.ReportMetric(float64(len(sBytes)-len(cmpBytes)), "BytesSaved")
	}
}

func BenchmarkStandardGzip(b *testing.B) {
	for i := 0; i < b.N; i++ {
		s := RandomString(1024)
		sBytes := []byte(s)
		t1 := time.Now()
		cmpBytes := standardCompress(sBytes, gzip.DefaultCompression)
		b.ReportMetric(float64(time.Since(t1).Nanoseconds()), "CompressTime")
		b.ReportMetric(float64(len(sBytes)-len(cmpBytes)), "BytesSaved")

		t1 = time.Now()
		_ = standardDecompress(cmpBytes)
		b.ReportMetric(float64(time.Since(t1).Nanoseconds()), "DecompressTime")
	}
}
<commit_msg>Update benchmark tests to use Compressor<commit_after>package network

import (
	"bytes"
	"compress/gzip"
	"fmt"
	"io/ioutil"
	"math/rand"
	"testing"
	"time"
)

func BenchmarkGzip(b *testing.B) {
	b.Run("standard gzip", BenchmarkStandardGzip)
	b.Run("best compression gzip", BenchmarkBestCompGzip)
	b.Run("best speed gzip", BenchmarkBestSpeedGzip)
}

func BenchmarkPooledVsNonPooledGzip(b *testing.B) {
	b.Run("PooledGzip", BenchmarkPooledGzip)
	b.Run("NonPooledGzip", BenchmarkStandardGzipCompression)
}

var (
	cmpBytes []byte
	pool     = NewCompressorPool()
)

func BenchmarkPooledGzip(b *testing.B) {
	for i := 0; i < b.N; i++ {
		s := RandomString(4096)
		sBytes := []byte(s)
		t1 := time.Now()
		w := pool.Get().(Compressor)
		cmpBytes, err := w.Compress(sBytes)
		if err != nil {
			b.Fatal(err)
		}
		w.Reset()
		pool.Put(w)
		b.ReportMetric(float64(time.Since(t1).Nanoseconds()), "CompressTime")
		b.ReportMetric(float64(len(sBytes)-len(cmpBytes)), "BytesSaved")
	}
}

func BenchmarkStandardGzipCompression(b *testing.B) {
	for i := 0; i < b.N; i++ {
		s := RandomString(4096)
		sBytes := []byte(s)
		t1 := time.Now()
		cmpBytes = standardCompress(sBytes, gzip.DefaultCompression)
		b.ReportMetric(float64(time.Since(t1).Nanoseconds()), "CompressTime")
		b.ReportMetric(float64(len(sBytes)-len(cmpBytes)), "BytesSaved")
	}
}

func BenchmarkStandardGzip(b *testing.B) {
	for i := 0; i < b.N; i++ {
		s := RandomString(1024)
		sBytes := []byte(s)
		t1 := time.Now()
		cmpBytes := standardCompress(sBytes, gzip.DefaultCompression)
		b.ReportMetric(float64(time.Since(t1).Nanoseconds()), "CompressTime")
		b.ReportMetric(float64(len(sBytes)-len(cmpBytes)), "BytesSaved")

		t1 = time.Now()
		_ = standardDecompress(cmpBytes)
		b.ReportMetric(float64(time.Since(t1).Nanoseconds()), "DecompressTime")
	}
}

{'fraction_non_alphanumeric': 0.16719473555934683, 'fraction_numerical': 0.018035583719229833, 'mean_word_length': 4.7478991596638656, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3701444', 'n_tokens_mistral': 1772, 'n_tokens_neox': 1625, 'n_words': 288}