From 7c99210e68f728c83f0078b37d76aab4b5cb0149 Mon Sep 17 00:00:00 2001
From: Will Norris <will@tailscale.com>
Date: Mon, 27 Mar 2023 14:34:56 -0700
Subject: [PATCH] log: allow toggling sockstat logs via c2n component logging

Signed-off-by: Will Norris <will@tailscale.com>
---
 ipn/ipnlocal/local.go     |  8 ++++++++
 log/sockstatlog/logger.go | 37 ++++++++++++++++++++++++++-----------
 2 files changed, 34 insertions(+), 11 deletions(-)

diff --git a/ipn/ipnlocal/local.go b/ipn/ipnlocal/local.go
index b27509a22..4c5efd9b7 100644
--- a/ipn/ipnlocal/local.go
+++ b/ipn/ipnlocal/local.go
@@ -315,6 +315,9 @@ func NewLocalBackend(logf logger.Logf, logID logid.PublicID, store ipn.StateStor
 		if err != nil {
 			log.Printf("error setting up sockstat logger: %v", err)
 		}
+		if b.sockstatLogger != nil {
+			b.sockstatLogger.SetLoggingEnabled(true)
+		}
 	}
 
 	// Default filter blocks everything and logs nothing, until Start() is called.
@@ -365,6 +368,7 @@ type componentLogState struct {
 
 var debuggableComponents = []string{
 	"magicsock",
+	"sockstats",
 }
 
 func componentStateKey(component string) ipn.StateKey {
@@ -389,6 +393,10 @@ func (b *LocalBackend) SetComponentDebugLogging(component string, until time.Tim
 			return err
 		}
 		setEnabled = mc.SetDebugLoggingEnabled
+	case "sockstats":
+		if b.sockstatLogger != nil {
+			setEnabled = b.sockstatLogger.SetLoggingEnabled
+		}
 	}
 	if setEnabled == nil || !slices.Contains(debuggableComponents, component) {
 		return fmt.Errorf("unknown component %q", component)
diff --git a/log/sockstatlog/logger.go b/log/sockstatlog/logger.go
index 744b8d630..1906e16e3 100644
--- a/log/sockstatlog/logger.go
+++ b/log/sockstatlog/logger.go
@@ -12,6 +12,7 @@
 	"net/http"
 	"os"
 	"path/filepath"
+	"sync/atomic"
 	"time"
 
 	"tailscale.com/logpolicy"
@@ -29,6 +30,8 @@
 
 // Logger logs statistics about network sockets.
 type Logger struct {
+	enabled atomic.Bool
+
 	ctx      context.Context
 	cancelFn context.CancelFunc
 
@@ -69,7 +72,7 @@ func SockstatLogID(logID logid.PublicID) logid.PrivateID {
 
 // NewLogger returns a new Logger that will store stats in logdir.
 // On platforms that do not support sockstat logging, a nil Logger will be returned.
-// The returned Logger must be shut down with Shutdown when it is no longer needed.
+// The returned Logger is not yet enabled, and must be shut down with Shutdown when it is no longer needed.
 func NewLogger(logdir string, logf logger.Logf, logID logid.PublicID) (*Logger, error) {
 	if !sockstats.IsAvailable {
 		return nil, nil
@@ -84,14 +87,11 @@ func NewLogger(logdir string, logf logger.Logf, logID logid.PublicID) (*Logger,
 		return nil, err
 	}
 
-	ctx, cancel := context.WithCancel(context.Background())
 	logger := &Logger{
-		ctx:      ctx,
-		cancelFn: cancel,
-		ticker:   time.NewTicker(pollPeriod),
-		logf:     logf,
-		filch:    filch,
-		tr:       logpolicy.NewLogtailTransport(logtail.DefaultHost),
+		ticker: time.NewTicker(pollPeriod),
+		logf:   logf,
+		filch:  filch,
+		tr:     logpolicy.NewLogtailTransport(logtail.DefaultHost),
 	}
 	logger.logger = logtail.NewLogger(logtail.Config{
 		BaseURL:    logpolicy.LogURL(),
@@ -114,11 +114,24 @@ func NewLogger(logdir string, logf logger.Logf, logID logid.PublicID) (*Logger,
 		HTTPC: &http.Client{Transport: logger.tr},
 	}, logf)
 
-	go logger.poll()
-
 	return logger, nil
 }
 
+// SetLoggingEnabled enables or disables logging.
+// When disabled, socket stats are not polled and no new logs are written to disk.
+// Existing logs can still be fetched via the C2N API.
+func (l *Logger) SetLoggingEnabled(v bool) {
+	old := l.enabled.Load()
+	if old != v && l.enabled.CompareAndSwap(old, v) {
+		if v {
+			l.ctx, l.cancelFn = context.WithCancel(context.Background())
+			go l.poll()
+		} else {
+			l.cancelFn()
+		}
+	}
+}
+
 func (l *Logger) Write(p []byte) (int, error) {
 	return l.logger.Write(p)
 }
@@ -173,7 +186,9 @@ func (l *Logger) Flush() {
 }
 
 func (l *Logger) Shutdown() {
-	l.cancelFn()
+	if l.cancelFn != nil {
+		l.cancelFn()
+	}
 	l.ticker.Stop()
 	l.filch.Close()
 	l.logger.Shutdown(context.Background())