internal/counter: don't write counters to disk if mode=off

It was decided in golang/go#63832 that when the telemetry mode is "off"
no telemetry data should be written to the file system. However, the
current implementation of "off" still creates the counter file -- it
just doesn't produce any reports.

Fix this to not even create the counter file when the mode is off. This
was rather tricky to implement, as it required auditing a lot of code to
see that we don't reach openMapped. However, per discussion with pjw@ it
should be sufficient to implement this check in Open, as is done here.

This broke some tests because
1. some tests were reading the wrong mode file (fixed by setting
   telemetry.ModeFile in counter_test.go)
2. the E2E tests were incorrectly escaping the RunProg test.run regexp
   (fixed by simplifying the RunProg logic)

For golang/go#63832

Change-Id: I47066a97a8fd17c4c2be776077d718e9cdbaf65a
Reviewed-on: https://go-review.googlesource.com/c/telemetry/+/542055
Auto-Submit: Robert Findley <[email protected]>
LUCI-TryBot-Result: Go LUCI <[email protected]>
Reviewed-by: Peter Weinberger <[email protected]>

Author: findleyr

counter/counter.go

@@ -67,10 +67,13 @@ func NewStack(name string, depth int) *StackCounter {
 	return counter.NewStack(name, depth)
 }
 
-// Open opens the counter file on disk and starts to mmap telemetry
-// counters to the file. Open also persists counters created and
-// incremented before it is called.
-// Programs are supposed to call this once.
+// Open prepares telemetry counters for recording to the file system.
+//
+// If the telemetry mode is "off", Open is a no-op. Otherwise, it opens the
+// counter file on disk and starts to mmap telemetry counters to the file.
+// Open also persists any counters already created in the current process.
+//
+// Programs using telemetry should call Open exactly once.
 func Open() {
 	counter.Open()
 }

counter/countertest/countertest.go

@@ -12,7 +12,6 @@ package countertest
 
 import (
 	"fmt"
-	"os"
 	"path/filepath"
 	"sync"
 	"testing"
@@ -51,11 +50,6 @@ func Open(telemetryDir string) {
 	telemetry.ModeFile = telemetry.ModeFilePath(filepath.Join(telemetryDir, "mode"))
 	telemetry.LocalDir = filepath.Join(telemetryDir, "local")
 	telemetry.UploadDir = filepath.Join(telemetryDir, "upload")
-	err1 := os.MkdirAll(telemetry.LocalDir, 0755)
-	err2 := os.MkdirAll(telemetry.UploadDir, 0755)
-	if err1 != nil || err2 != nil {
-		panic(fmt.Sprintf("failed to create telemetry dirs: %v, %v", err1, err2))
-	}
 
 	counter.Open()
 	opened = true

internal/counter/counter_test.go

@@ -101,6 +101,9 @@ func TestParallel(t *testing.T) {
 		t.Fatal(f.err)
 	}
 	current := f.current.Load()
+	if current == nil {
+		t.Fatal("no mapped file")
+	}
 	name := current.f.Name()
 	t.Logf("wrote %s:\n%s", name, hexDump(current.mapping.Data))
 
@@ -499,7 +502,11 @@ func TestStack(t *testing.T) {
 		}
 	}
 	// check that Parse expands compressed counter names
-	data := f.current.Load().mapping.Data
+	current := f.current.Load()
+	if current == nil {
+		t.Fatal("no mapped file")
+	}
+	data := current.mapping.Data
 	fname := "2023-01-01.v1.count" // bogus file name required by Parse.
 	theFile, err := Parse(fname, data)
 	if err != nil {
@@ -547,8 +554,8 @@ func setup(t *testing.T) {
 	telemetry.UploadDir = tmpDir + "/upload"
 	os.MkdirAll(telemetry.LocalDir, 0777)
 	os.MkdirAll(telemetry.UploadDir, 0777)
+	telemetry.ModeFile = telemetry.ModeFilePath(filepath.Join(tmpDir, "mode"))
 	// os.UserConfigDir() is "" in tests so no point in looking at it
-
 }
 
 func restore() {

internal/counter/file.go

@@ -38,7 +38,7 @@ type file struct {
 	namePrefix string
 	err        error
 	meta       string
-	current    atomic.Pointer[mappedFile] // can be read without holding mu
+	current    atomic.Pointer[mappedFile] // can be read without holding mu, but may be nil
 }
 
 var defaultFile file
@@ -350,6 +350,11 @@ var mainCounter = New("counter/main")
 // any reports are generated.
 // (Otherwise expired count files will not be deleted on Windows.)
 func Open() func() {
+	if mode, _ := telemetry.Mode(); mode == "off" {
+		// Don't open the file when telemetry is off.
+		defaultFile.err = ErrDisabled
+		return func() {} // No need to clean up.
+	}
 	debugPrintf("Open")
 	mainCounter.Add(1)
 	defaultFile.rotate()

internal/regtest/e2e_test.go

@@ -21,47 +21,98 @@ import (
 	"golang.org/x/telemetry/counter"
 	"golang.org/x/telemetry/internal/config"
 	icounter "golang.org/x/telemetry/internal/counter"
+	it "golang.org/x/telemetry/internal/telemetry"
 	"golang.org/x/telemetry/internal/testenv"
 )
 
+// programs to run
+const (
+	progIncCounters = "inccounters"
+	prog1           = "prog1"
+	prog2           = "prog2"
+)
+
+func TestMain(m *testing.M) {
+	Main(m, map[string]func() int{
+		progIncCounters: func() int {
+			counter.Inc("counter")
+			counter.Inc("counter:surprise")
+			counter.New("gopls/editor:expected").Inc()
+			counter.New("gopls/editor:surprise").Inc()
+			counter.NewStack("stack/expected", 1).Inc()
+			counter.NewStack("stack-surprise", 1).Inc()
+			return 0
+		},
+		prog1: func() int {
+			fmt.Println("FuncB")
+			return 0
+		},
+		prog2: func() int {
+			fmt.Println("FuncC")
+			return 1
+		},
+	})
+}
+
 func TestRunProg(t *testing.T) {
 	testenv.MustHaveExec(t)
 	telemetryDir := t.TempDir()
-	fmt.Println("START")
 	t.Run("prog1", func(t *testing.T) {
-		if out, err := RunProg(t, telemetryDir, func() int {
-			fmt.Println("FuncB")
-			return 0
-		}); err != nil || !bytes.Contains(out, []byte("START")) || !bytes.Contains(out, []byte("FuncB")) {
-			t.Errorf("first RunProg = (%s, %v), want 'START', 'FuncB', and succeed", out, err)
+		if out, err := RunProg(telemetryDir, prog1); err != nil || !bytes.Contains(out, []byte("FuncB")) || bytes.Contains(out, []byte("FuncC")) {
+			t.Errorf("first RunProg = (%s, %v), want FuncB' and succeed", out, err)
 		}
 	})
-	fmt.Println("MIDDLE")
 	t.Run("prog2", func(t *testing.T) {
-		if out, err := RunProg(t, telemetryDir, func() int {
-			fmt.Println("FuncC")
-			return 1
-		}); err == nil || !bytes.Contains(out, []byte("START")) || bytes.Contains(out, []byte("FuncB")) || !bytes.Contains(out, []byte("MIDDLE")) || !bytes.Contains(out, []byte("FuncC")) {
-			t.Errorf("second RunProg = (%s, %v), want 'START', 'MIDDLE', 'FuncC' (but no 'FuncB') and fail", out, err)
+		if out, err := RunProg(telemetryDir, prog2); err == nil || bytes.Contains(out, []byte("FuncB")) || !bytes.Contains(out, []byte("FuncC")) {
+			t.Errorf("second RunProg = (%s, %v), want 'FuncC' and fail", out, err)
 		}
 	})
 }
 
+func TestE2E_off(t *testing.T) {
+	testenv.MustHaveExec(t)
+
+	tests := []struct {
+		mode         string // if empty, don't set the mode
+		wantLocalDir bool
+	}{
+		{"", true},
+		{"local", true},
+		{"on", true},
+		{"off", false},
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("mode=%s", test.mode), func(t *testing.T) {
+			telemetryDir := t.TempDir()
+			if test.mode != "" {
+				if err := it.ModeFilePath(filepath.Join(telemetryDir, "mode")).SetMode(test.mode); err != nil {
+					t.Fatalf("SetMode failed: %v", err)
+				}
+			}
+			out, err := RunProg(telemetryDir, progIncCounters)
+			if err != nil {
+				t.Fatalf("program failed unexpectedly (%v)\n%s", err, out)
+			}
+			localDir := filepath.Join(telemetryDir, "local")
+			_, err = os.Stat(localDir)
+			if err != nil && !os.IsNotExist(err) {
+				t.Fatalf("os.Stat(%q): unexpected error: %v", localDir, err)
+			}
+			if gotLocalDir := err == nil; gotLocalDir != test.wantLocalDir {
+				t.Errorf("got /local dir: %v, want %v; out:\n%s", gotLocalDir, test.wantLocalDir, string(out))
+			}
+		})
+	}
+}
+
 func TestE2E(t *testing.T) {
 	testenv.MustHaveExec(t)
 	telemetryDir := t.TempDir()
 
 	goVers, progVers, progName := ProgInfo(t)
 
-	out, err := RunProg(t, telemetryDir, func() int {
-		counter.Inc("counter")
-		counter.Inc("counter:surprise")
-		counter.New("gopls/editor:expected").Inc()
-		counter.New("gopls/editor:surprise").Inc()
-		counter.NewStack("stack/expected", 1).Inc()
-		counter.NewStack("stack-surprise", 1).Inc()
-		return 0
-	})
+	out, err := RunProg(telemetryDir, progIncCounters)
 	if err != nil {
 		t.Fatalf("program failed unexpectedly (%v)\n%s", err, out)
 	}

internal/regtest/regtest.go

@@ -9,67 +9,57 @@
 package regtest
 
 import (
+	"flag"
+	"fmt"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"runtime/debug"
 	"strings"
-	"sync"
 	"testing"
 
 	"golang.org/x/telemetry/counter/countertest"
 )
 
-const telemetryDirEnvVar = "_COUNTERTEST_RUN_TELEMETRY_DIR"
-
-var (
-	runProgMu  sync.Mutex
-	hasRunProg = map[string]bool{} // test name -> RunProg already called for this test.
+const (
+	telemetryDirEnvVar = "_COUNTERTEST_RUN_TELEMETRY_DIR"
+	entryPointEnvVar   = "_COUNTERTEST_ENTRYPOINT"
 )
 
-func canRunProg(name string) bool {
-	runProgMu.Lock()
-	defer runProgMu.Unlock()
-	if hasRunProg[name] {
-		return false
-	}
-	hasRunProg[name] = true
-	return true
-}
-
-// RunProg runs prog in a separate process with the specified telemetry directory.
-// The return value of prog is the exit code of the process.
-// RunProg can be called at most once per testing.T.
-// If an integration test needs to run multiple programs, use subtests.
-func RunProg(t *testing.T, telemetryDir string, prog func() int) ([]byte, error) {
-	testName := t.Name()
-	if !canRunProg(testName) {
-		t.Fatalf("RunProg was called more than once in test %v. Use subtests if a test needs it more than once", testName)
+// Main is a test main function for use in TestMain, which runs one of the
+// given programs when invoked as a separate process via RunProg.
+//
+// The return value of each program is the exit code of the process.
+func Main(m *testing.M, programs map[string]func() int) {
+	if d := os.Getenv(telemetryDirEnvVar); d != "" {
+		countertest.Open(d)
 	}
-	if telemetryDir := os.Getenv(telemetryDirEnvVar); telemetryDir != "" {
-		// run the prog.
-		countertest.Open(telemetryDir)
-		os.Exit(prog())
+	if e, ok := os.LookupEnv(entryPointEnvVar); ok {
+		if prog, ok := programs[e]; ok {
+			os.Exit(prog())
+		}
+		fmt.Fprintf(os.Stderr, "unknown program %q", e)
+		os.Exit(2)
 	}
+	flag.Parse()
+	os.Exit(m.Run())
+}
 
-	t.Helper()
-
+// RunProg runs the named program in a separate process with the specified
+// telemetry directory, where prog is one of the programs passed to Main (which
+// must be invoked by TestMain).
+func RunProg(telemetryDir string, prog string) ([]byte, error) {
 	testBin, err := os.Executable()
 	if err != nil {
-		t.Fatalf("cannot determine the current process's executable name: %v", err)
+		return nil, fmt.Errorf("cannot determine the current process's executable name: %v", err)
 	}
 
 	// Spawn a subprocess to run the `prog`, by setting subprocessKeyEnvVar and telemetryDirEnvVar.
-	cmd := exec.Command(testBin, "-test.run", testName)
-	cmd.Env = append(cmd.Env, telemetryDirEnvVar+"="+telemetryDir)
+	cmd := exec.Command(testBin)
+	cmd.Env = append(cmd.Env, telemetryDirEnvVar+"="+telemetryDir, entryPointEnvVar+"="+prog)
 	return cmd.CombinedOutput()
 }
 
-// InSubprocess returns whether the current process is a subprocess forked by RunProg.
-func InSubprocess() bool {
-	return os.Getenv(telemetryDirEnvVar) != ""
-}
-
 // ProgInfo returns the go version, program name and version info the process would record in its counter file.
 func ProgInfo(t *testing.T) (goVersion, progVersion, progName string) {
 	info, ok := debug.ReadBuildInfo()