From ceaaa239620e956a45cb194b26fc2a9afad9eba9 Mon Sep 17 00:00:00 2001
From: Josh Bleecher Snyder <josh@tailscale.com>
Date: Mon, 24 May 2021 17:00:46 -0700
Subject: [PATCH] wgengine/wglog: cache strings
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

We repeat many peers each time we call SetPeers.
Instead of constructing strings for them from scratch every time,
keep strings alive across iterations.

name        old time/op    new time/op    delta
SetPeers-8    3.58µs ± 1%    2.41µs ± 1%  -32.60%  (p=0.000 n=9+10)

name        old alloc/op   new alloc/op   delta
SetPeers-8    2.53kB ± 0%    1.30kB ± 0%  -48.73%  (p=0.000 n=10+10)

name        old allocs/op  new allocs/op  delta
SetPeers-8      99.0 ± 0%      16.0 ± 0%  -83.84%  (p=0.000 n=10+10)

We could reduce alloc/op 12% and allocs/op 23% if strs had
type map[string]strCache instead of map[string]*strCache,
but that wipes out the execution time impact.
Given that re-use is the most common scenario, let's optimize for it.

Signed-off-by: Josh Bleecher Snyder <josh@tailscale.com>
---
 wgengine/wglog/wglog.go | 35 +++++++++++++++++++++++++++++++----
 1 file changed, 31 insertions(+), 4 deletions(-)

diff --git a/wgengine/wglog/wglog.go b/wgengine/wglog/wglog.go
index 4a9c5a6a9..59b740ec1 100644
--- a/wgengine/wglog/wglog.go
+++ b/wgengine/wglog/wglog.go
@@ -9,6 +9,7 @@
 	"encoding/base64"
 	"fmt"
 	"strings"
+	"sync"
 	"sync/atomic"
 
 	"github.com/tailscale/wireguard-go/device"
@@ -21,7 +22,15 @@
 // It can be modified at run time to adjust to new wireguard-go configurations.
 type Logger struct {
 	DeviceLogger *device.Logger
-	replace      atomic.Value // of map[string]string
+	replace      atomic.Value            // of map[string]string
+	mu           sync.Mutex              // protects strs
+	strs         map[wgkey.Key]*strCache // cached strs used to populate replace
+}
+
+// strCache holds a wireguard-go and a Tailscale style peer string.
+type strCache struct {
+	wg, ts string
+	used   bool // track whether this strCache was used in a particular round
 }
 
 // NewLogger creates a new logger for use with wireguard-go.
@@ -76,18 +85,36 @@ func NewLogger(logf logger.Logf) *Logger {
 		Verbosef: logger.WithPrefix(wrapper, "[v2] "),
 		Errorf:   wrapper,
 	}
+	ret.strs = make(map[wgkey.Key]*strCache)
 	return ret
 }
 
 // SetPeers adjusts x to rewrite the peer public keys found in peers.
 // SetPeers is safe for concurrent use.
 func (x *Logger) SetPeers(peers []wgcfg.Peer) {
+	x.mu.Lock()
+	defer x.mu.Unlock()
 	// Construct a new peer public key log rewriter.
 	replace := make(map[string]string)
 	for _, peer := range peers {
-		old := wireguardGoString(peer.PublicKey)
-		new := peer.PublicKey.ShortString()
-		replace[old] = new
+		c, ok := x.strs[peer.PublicKey] // look up cached strs
+		if !ok {
+			wg := wireguardGoString(peer.PublicKey)
+			ts := peer.PublicKey.ShortString()
+			c = &strCache{wg: wg, ts: ts}
+			x.strs[peer.PublicKey] = c
+		}
+		c.used = true
+		replace[c.wg] = c.ts
+	}
+	// Remove any unused cached strs.
+	for k, c := range x.strs {
+		if !c.used {
+			delete(x.strs, k)
+			continue
+		}
+		// Mark c as unused for next round.
+		c.used = false
 	}
 	x.replace.Store(replace)
 }