use smallzstd and sync pool

This commit is contained in:
Johan Siebens 2023-01-02 07:48:30 +01:00 committed by Juan Font
parent 6d669c6b9c
commit e128796b59
3 changed files with 30 additions and 4 deletions

View File

@ -33,7 +33,7 @@
# When updating go.mod or go.sum, a new sha will need to be calculated,
# update this if you have a mismatch after doing a change to thos files.
vendorSha256 = "sha256-SuKT+b8g6xEK15ry2IAmpS/vwDG+zJqK9nfsWpHNXuU=";
vendorSha256 = "sha256-8p5NFxXKaZPsW4B6NMzfi0pqfVroIahSgA0fukvB3JI=";
ldflags = ["-s" "-w" "-X github.com/juanfont/headscale/cmd/headscale/cli.Version=v${version}"];
};

2
go.mod
View File

@ -7,6 +7,7 @@ require (
github.com/ccding/go-stun/stun v0.0.0-20200514191101-4dc67bcdb029
github.com/cenkalti/backoff/v4 v4.2.0
github.com/coreos/go-oidc/v3 v3.4.0
github.com/davecgh/go-spew v1.1.1
github.com/deckarep/golang-set/v2 v2.1.0
github.com/efekarakus/termcolor v1.0.1
github.com/glebarez/sqlite v1.5.0
@ -58,7 +59,6 @@ require (
github.com/cespare/xxhash/v2 v2.2.0 // indirect
github.com/containerd/console v1.0.3 // indirect
github.com/containerd/continuity v0.3.0 // indirect
github.com/davecgh/go-spew v1.1.1 // indirect
github.com/docker/cli v20.10.21+incompatible // indirect
github.com/docker/docker v20.10.21+incompatible // indirect
github.com/docker/go-connections v0.4.0 // indirect

View File

@ -3,9 +3,11 @@ package headscale
import (
"encoding/binary"
"encoding/json"
"sync"
"github.com/klauspost/compress/zstd"
"github.com/rs/zerolog/log"
"tailscale.com/smallzstd"
"tailscale.com/tailcfg"
"tailscale.com/types/key"
)
@ -103,8 +105,7 @@ func (h *Headscale) marshalMapResponse(
var respBody []byte
if compression == ZstdCompression {
encoder, _ := zstd.NewWriter(nil)
respBody = encoder.EncodeAll(jsonBody, nil)
respBody = zstdEncode(jsonBody)
if !isNoise { // if legacy protocol
respBody = h.privateKey.SealTo(machineKey, respBody)
}
@ -122,3 +123,28 @@ func (h *Headscale) marshalMapResponse(
return data, nil
}
func zstdEncode(in []byte) []byte {
encoder, ok := zstdEncoderPool.Get().(*zstd.Encoder)
if !ok {
panic("invalid type in sync pool")
}
out := encoder.EncodeAll(in, nil)
_ = encoder.Close()
zstdEncoderPool.Put(encoder)
return out
}
var zstdEncoderPool = &sync.Pool{
New: func() any {
encoder, err := smallzstd.NewEncoder(
nil,
zstd.WithEncoderLevel(zstd.SpeedFastest))
if err != nil {
panic(err)
}
return encoder
},
}