From 703d78900582dc0c77da1895c34cb7b75a93c45e Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Tue, 7 Apr 2020 22:24:06 -0700
Subject: [PATCH] tailcfg: add MapResponse.Debug mechanism to trigger logging
 heap pprof

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 control/controlclient/direct.go |  4 +++
 log/logheap/logheap.go          | 45 +++++++++++++++++++++++++++++++++
 log/logheap/logheap_test.go     | 40 +++++++++++++++++++++++++++++
 tailcfg/tailcfg.go              | 13 ++++++++++
 4 files changed, 102 insertions(+)
 create mode 100644 log/logheap/logheap.go
 create mode 100644 log/logheap/logheap_test.go

diff --git a/control/controlclient/direct.go b/control/controlclient/direct.go
index 9078e9374..626afab87 100644
--- a/control/controlclient/direct.go
+++ b/control/controlclient/direct.go
@@ -27,6 +27,7 @@ import (
 	"github.com/tailscale/wireguard-go/wgcfg"
 	"golang.org/x/crypto/nacl/box"
 	"golang.org/x/oauth2"
+	"tailscale.com/log/logheap"
 	"tailscale.com/net/tlsdial"
 	"tailscale.com/tailcfg"
 	"tailscale.com/types/logger"
@@ -588,6 +589,9 @@ func (c *Direct) PollNetMap(ctx context.Context, maxPolls int, cb func(*NetworkM
 			vlogf("netmap: new map contains DERP map")
 			lastDERPMap = resp.DERPMap
 		}
+		if resp.Debug != nil && resp.Debug.LogHeapPprof {
+			logheap.LogHeap()
+		}
 
 		nm := &NetworkMap{
 			NodeKey:      tailcfg.NodeKey(persist.PrivateNodeKey.Public()),
diff --git a/log/logheap/logheap.go b/log/logheap/logheap.go
new file mode 100644
index 000000000..a97d7f6b9
--- /dev/null
+++ b/log/logheap/logheap.go
@@ -0,0 +1,45 @@
+// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+// Package logheap logs a heap pprof profile.
+package logheap
+
+import (
+	"bytes"
+	"encoding/json"
+	"io"
+	"os"
+	"runtime"
+	"runtime/pprof"
+	"time"
+)
+
+// LogHeap writes a JSON logtail record with the base64 heap pprof to
+// os.Stderr.
+func LogHeap() {
+	logHeap(os.Stderr)
+}
+
+type logTail struct {
+	ClientTime string `json:"client_time"`
+}
+
+type pprofRec struct {
+	Heap []byte `json:"heap,omitempty"`
+}
+
+type logLine struct {
+	LogTail logTail  `json:"logtail"`
+	Pprof   pprofRec `json:"pprof"`
+}
+
+func logHeap(w io.Writer) error {
+	runtime.GC()
+	buf := new(bytes.Buffer)
+	pprof.WriteHeapProfile(buf)
+	return json.NewEncoder(w).Encode(logLine{
+		LogTail: logTail{ClientTime: time.Now().Format(time.RFC3339Nano)},
+		Pprof:   pprofRec{Heap: buf.Bytes()},
+	})
+}
diff --git a/log/logheap/logheap_test.go b/log/logheap/logheap_test.go
new file mode 100644
index 000000000..d55cb5d35
--- /dev/null
+++ b/log/logheap/logheap_test.go
@@ -0,0 +1,40 @@
+// Copyright (c) 2020 Tailscale Inc & AUTHORS All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package logheap
+
+import (
+	"bytes"
+	"compress/gzip"
+	"encoding/json"
+	"io/ioutil"
+	"testing"
+)
+
+func TestLogHeap(t *testing.T) {
+	var buf bytes.Buffer
+	if err := logHeap(&buf); err != nil {
+		t.Fatal(err)
+	}
+	t.Logf("Got line: %s", buf.Bytes())
+
+	var ll logLine
+	if err := json.Unmarshal(buf.Bytes(), &ll); err != nil {
+		t.Fatal(err)
+	}
+
+	zr, err := gzip.NewReader(bytes.NewReader(ll.Pprof.Heap))
+	if err != nil {
+		t.Fatal(err)
+	}
+	rawProto, err := ioutil.ReadAll(zr)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// Just sanity check it. Too lazy to properly decode the protobuf. But see that
+	// it contains an expected sample name.
+	if !bytes.Contains(rawProto, []byte("alloc_objects")) {
+		t.Errorf("raw proto didn't contain `alloc_objects`: %q", rawProto)
+	}
+}
diff --git a/tailcfg/tailcfg.go b/tailcfg/tailcfg.go
index 5afe05129..83a148b47 100644
--- a/tailcfg/tailcfg.go
+++ b/tailcfg/tailcfg.go
@@ -500,6 +500,19 @@ type MapResponse struct {
 	Roles        []Role
 	// TODO: Groups       []Group
 	// TODO: Capabilities []Capability
+
+	// Debug is normally nil, except for when the control server
+	// is setting debug settings on a node.
+	Debug *Debug `json:",omitempty"`
+}
+
+// Debug are instructions from the control server to the client
+// to adjust debug settings.
+type Debug struct {
+	// LogHeapPprof controls whether the client should logs
+	// its heap pprof data. Each true value sent from the server
+	// means that client should do one more log.
+	LogHeapPprof bool `json:",omitempty"`
 }
 
 func (k MachineKey) String() string { return fmt.Sprintf("mkey:%x", k[:]) }