Just a tool to bench compression algorithms for your dataset

compressor.go

package main

import (
	"flag"
	"fmt"
	"io"
	"net/http"
	"time"

	gzip "github.com/klauspost/compress/gzip"
	"github.com/klauspost/compress/s2"
	"github.com/klauspost/compress/snappy"
	"github.com/klauspost/compress/zlib"
	"github.com/klauspost/compress/zstd"
	"github.com/pierrec/lz4"

	"github.com/gin-gonic/gin"
	log "github.com/sirupsen/logrus"
	"github.com/thanhpk/randstr"
	"github.com/urfave/negroni"

	"github.com/miolini/datacounter"
)

type App struct {
	Server       *http.Server
	ListenAddr   string
	WriteTimeout int
	ReadTimeout  int
}

type CompressArgs struct {
	Algorithm string `form:"algo"`
}

func Compress(c *gin.Context) {
	var r CompressArgs
	err := c.ShouldBindQuery(&r)
	if err != nil {
		c.JSON(400, gin.H{"error": "ARGUMENTS", "message": "bindig failed: " + err.Error()})
	}

	dc := datacounter.NewWriterCounter(io.Discard)
	var w io.Writer

	start := time.Now()

	switch r.Algorithm {
	case "gzip":
		w = gzip.NewWriter(dc)
	case "zstd":
		w, err = zstd.NewWriter(dc)
		if err != nil {
			log.Errorf("creating zstd failed: %s", err)
		}
	case "s2":
		w = s2.NewWriter(dc)
	case "snappy":
		w = snappy.NewBufferedWriter(dc)
	case "zlib":
		w = zlib.NewWriter(dc)
	case "lz4":
		w = lz4.NewWriter(dc)
	default:
		w = dc
	}

	if w == nil {
		c.JSON(500, gin.H{"error": "INTENAL", "message": "smthing bad happened"})
		return
	}

	nw, err := io.Copy(w, c.Request.Body)

	if err != nil {
		c.JSONP(500, gin.H{"error": "pipe", "message": err.Error()})
	}

	if w != dc {
		w.(io.WriteCloser).Close()
	}

	dur := time.Since(start)
	com := (1 - float64(dc.Count())/float64(nw)) * 100

	c.JSONP(200, gin.H{
		"read":  nw,
		"write": dc.Count(),
		"T":     dur.Truncate(time.Millisecond).String(),
		"All":   time.Since(c.Keys["started"].(time.Time)).Truncate(time.Millisecond).String(),
		"%":     com,
		"s":     fmt.Sprintf("%.3fMB/s", float64(nw)/1024/1024/dur.Seconds()),
	})
}

func EndLogMiddleware(w http.ResponseWriter, r *http.Request, next http.HandlerFunc) {
	nw := w.(negroni.ResponseWriter)
	start := time.Now()

	rid := r.Header.Get("X-Request-ID")
	if rid == "" {
		rid = randstr.String(8)
		r.Header.Set("X-Request-ID", rid)
	}

	nw.Header().Set("X-Request-ID", rid)
	nw.Header().Set("X-Started", fmt.Sprintf("%d", start.Unix()))
	nw.Header().Set("X-Server", "Compressor")

	next(w, r)

	log.WithFields(log.Fields{
		"method": r.Method,
		"path":   r.URL.Path,
		"CL":     r.ContentLength,
		"UA":     r.Header.Get("User-Agent"),
		"XRI":    r.Header.Get("X-Real-IP"),
		"XFF":    r.Header.Get("X-Forwarded-From"),
		"RID":    r.Header.Get("X-Request-ID"),
		"IP":     r.RemoteAddr,
		"T":      time.Since(start).Round(10 * time.Microsecond).String(),
		"RES":    nw.Size(),
	}).Infof("[END=%d]", nw.Status())
}

var app App = App{}

func init() {
	flag.StringVar(&app.ListenAddr, "l", ":8080", "listen addr")
	flag.IntVar(&app.WriteTimeout, "w", 15, "write timeout in seconds")
	flag.IntVar(&app.ReadTimeout, "r", 15, "read timeout in seconds")
	flag.Parse()
}

func main() {
	// w := gzip.NewWriter(io.Discard)
	gin.SetMode(gin.ReleaseMode)
	gin.DisableConsoleColor()

	n := negroni.New()
	n.Use(negroni.HandlerFunc(EndLogMiddleware))

	r := gin.Default()
	r.Use(func(ctx *gin.Context) {
		ctx.Keys = make(map[string]interface{})
		ctx.Keys["started"] = time.Now()
	})

	r.PUT("/compress", Compress)
	n.UseHandler(r)

	app.Server = &http.Server{
		Handler:      n,
		Addr:         app.ListenAddr,
		WriteTimeout: time.Duration(app.WriteTimeout) * time.Second,
		ReadTimeout:  time.Duration(app.ReadTimeout) * time.Second,
	}

	log.Infof("Starting server on %s", app.Server.Addr)
	log.Fatal(app.Server.ListenAndServe())
}