WIP: compression

go.mod

@@ -11,6 +11,7 @@ require (
 	github.com/facebookgo/grace v0.0.0-20180706040059-75cf19382434
 	github.com/facebookgo/pidfile v0.0.0-20150612191647-f242e2999868
 	github.com/google/go-cmp v0.6.0
+	github.com/klauspost/compress v1.17.9
 	github.com/libp2p/go-reuseport v0.4.0
 	github.com/pkg/errors v0.9.1
 	github.com/prometheus/client_golang v1.20.2
@@ -68,7 +69,6 @@ require (
 	github.com/google/uuid v1.6.0 // indirect
 	github.com/josharian/intern v1.0.0 // indirect
 	github.com/json-iterator/go v1.1.12 // indirect
-	github.com/klauspost/compress v1.17.9 // indirect
 	github.com/kylelemons/godebug v1.1.0 // indirect
 	github.com/mailru/easyjson v0.7.7 // indirect
 	github.com/moby/docker-image-spec v1.3.1 // indirect

pkg/conf/clusters.go

@@ -54,6 +54,8 @@ type Host struct {
 	GRPC bool `toml:"grpc"`
 	// Optional, number of TCP connections to open to the host.
 	MTCP int `toml:"mtcp"`
+	// Optional, compressor to use, default - no compression
+	Compression string `toml:"compression"`
 }
 
 // ReadClustersConfig reads clusters set from a reader.

pkg/target/hostMTCP.go

@@ -3,6 +3,9 @@ package target
 import (
 	"bufio"
 	"fmt"
+	"github.com/klauspost/compress/gzip"
+	"github.com/klauspost/compress/s2"
+	"github.com/klauspost/compress/snappy"
 	"math/rand"
 	"net"
 	"strconv"
@@ -20,12 +23,13 @@ import (
 
 // HostMTCP represents a single target hosts to send records in multimple TCP connections to.
 type HostMTCP struct {
-	Name      string
-	Port      uint16
-	Ch        chan *rec.RecBytes
-	Available atomic.Bool
-	NumMTCP   int
-	MTCPs     []MultiConnection
+	Name        string
+	Port        uint16
+	Ch          chan *rec.RecBytes
+	Available   atomic.Bool
+	NumMTCP     int
+	MTCPs       []MultiConnection
+	Compression string
 
 	stop chan int
 
@@ -58,12 +62,25 @@ type MultiConnection struct {
 	Available   atomic.Bool
 }
 
+func newCompressedWriter(c net.Conn, bufSize int, typ string) *bufio.Writer {
+	switch typ {
+	case "s2":
+		return bufio.NewWriterSize(s2.NewWriter(c), bufSize)
+	case "snappy":
+		return bufio.NewWriterSize(snappy.NewBufferedWriter(c), bufSize)
+	case "gzip":
+		return bufio.NewWriterSize(gzip.NewWriter(c), bufSize)
+	default:
+		return bufio.NewWriterSize(c, bufSize)
+	}
+}
+
 // New or updated target connection from existing net.Conn
 // Requires Connection.Mutex lock
-func (c *MultiConnection) New(n net.Conn, bufSize int) {
+func (c *MultiConnection) New(n net.Conn, bufSize int, compressor string) {
 	c.Conn = n
 	c.LastConnUse = time.Now()
-	c.W = bufio.NewWriterSize(n, bufSize)
+	c.W = newCompressedWriter(n, bufSize, compressor)
 }
 
 // Close the connection while mainaining correct internal state.
@@ -123,6 +140,7 @@ func ConstructHostMTCP(clusterName string, mainCfg conf.Main, hostCfg conf.Host,
 	}
 	h.NumMTCP = hostCfg.MTCP
 	h.MTCPs = make([]MultiConnection, h.NumMTCP)
+	h.Compression = hostCfg.Compression
 
 	h.rnd = rand.New(rand.NewSource(time.Now().UnixNano()))
 	h.Ms = ms
@@ -338,7 +356,7 @@ func (h *HostMTCP) connect(c *MultiConnection, attemptCount int) {
 		}
 		return
 	}
-	c.New(conn, h.conf.TCPOutBufSize)
+	c.New(conn, h.conf.TCPOutBufSize, h.Compression)
 	c.Available.Store(true)
 	h.setHostAvailability()
 }

vendor/modules.txt

@@ -293,10 +293,15 @@ github.com/json-iterator/go
 # github.com/klauspost/compress v1.17.9
 ## explicit; go 1.20
 github.com/klauspost/compress
+github.com/klauspost/compress/flate
 github.com/klauspost/compress/fse
+github.com/klauspost/compress/gzip
 github.com/klauspost/compress/huff0
 github.com/klauspost/compress/internal/cpuinfo
+github.com/klauspost/compress/internal/race
 github.com/klauspost/compress/internal/snapref
+github.com/klauspost/compress/s2
+github.com/klauspost/compress/snappy
 github.com/klauspost/compress/zstd
 github.com/klauspost/compress/zstd/internal/xxhash
 # github.com/kylelemons/godebug v1.1.0