From e128796b59a6416ba75a385596519d519b4f0e4e Mon Sep 17 00:00:00 2001
From: Johan Siebens <johan.siebens@gmail.com>
Date: Mon, 2 Jan 2023 07:48:30 +0100
Subject: [PATCH] use smallzstd and sync pool

---
 flake.nix                |  2 +-
 go.mod                   |  2 +-
 protocol_common_utils.go | 30 ++++++++++++++++++++++++++++--
 3 files changed, 30 insertions(+), 4 deletions(-)

diff --git a/flake.nix b/flake.nix
index 01677979..8a3e59f2 100644
--- a/flake.nix
+++ b/flake.nix
@@ -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}"];
         };
diff --git a/go.mod b/go.mod
index cba5e99f..f35f19c9 100644
--- a/go.mod
+++ b/go.mod
@@ -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
diff --git a/protocol_common_utils.go b/protocol_common_utils.go
index d3fb6979..96c236d0 100644
--- a/protocol_common_utils.go
+++ b/protocol_common_utils.go
@@ -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
+	},
+}