From 3f84ba6b05e7234d5f0b5e66050b1314767c87bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Daniel=20Mart=C3=AD?= <mvdan@mvdan.cc>
Date: Fri, 26 Jul 2024 16:29:36 +0100
Subject: [PATCH] cmd/cue: do not initialize cmd.Command twice in `cue cmd`
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

This allows `cue cmd --cpuprofile` to work, avoids other issues
such as initializing cueexperiment and cuedebug twice unnecessarily,
and also adds an assertion so that future double-init bugs like this one
don't go unnoticed for a while.

Signed-off-by: Daniel Martí <mvdan@mvdan.cc>
Change-Id: I97dd014cb2997f20122dd761822a32e483d9e951
Reviewed-on: https://review.gerrithub.io/c/cue-lang/cue/+/1198496
TryBot-Result: CUEcueckoo <cueckoo@cuelang.org>
Unity-Result: CUE porcuepine <cue.porcuepine@gmail.com>
Reviewed-by: Matthew Sackman <matthew@cue.works>
---
 cmd/cue/cmd/custom.go                   |  8 +++++---
 cmd/cue/cmd/root.go                     | 10 ++++++++++
 cmd/cue/cmd/testdata/script/pprof.txtar |  4 +---
 3 files changed, 16 insertions(+), 6 deletions(-)

diff --git a/cmd/cue/cmd/custom.go b/cmd/cue/cmd/custom.go
index 17b8ed1b833..52bac0e9cce 100644
--- a/cmd/cue/cmd/custom.go
+++ b/cmd/cue/cmd/custom.go
@@ -136,13 +136,15 @@ func customCommand(c *Command, typ, name string, tools *cue.Instance) (*cobra.Co
 		Use:   usage,
 		Short: lookupString(o, "$short", short),
 		Long:  lookupString(o, "$long", long),
-		RunE: mkRunE(c, func(cmd *Command, args []string) error {
+		// Note that we don't use mkRunE here, as the parent func is already wrapped by
+		// another mkRunE call, and all Command initialization has already happened.
+		RunE: func(cmd *cobra.Command, args []string) error {
 			// TODO:
 			// - parse flags and env vars
 			// - constrain current config with config section
 
-			return doTasks(cmd, name, tools)
-		}),
+			return doTasks(c, name, tools)
+		},
 	}
 
 	// TODO: implement var/flag handling.
diff --git a/cmd/cue/cmd/root.go b/cmd/cue/cmd/root.go
index fe2028efbaf..11bea7b1723 100644
--- a/cmd/cue/cmd/root.go
+++ b/cmd/cue/cmd/root.go
@@ -85,8 +85,18 @@ type Stats struct {
 	}
 }
 
+var hasRunCommand bool
+
 func mkRunE(c *Command, f runFunction) func(*cobra.Command, []string) error {
 	return func(cmd *cobra.Command, args []string) error {
+		// The init work below should only happen once per cmd/cue invocation;
+		// if it happens twice, we'll misbehave by writing stats twice
+		// or miscalculating pprof and stats numbers.
+		if hasRunCommand {
+			panic("cmd/cue/cmd.mkRunE init ran twice")
+		}
+		hasRunCommand = true
+
 		c.Command = cmd
 
 		statsEnc, err := statsEncoder(c)
diff --git a/cmd/cue/cmd/testdata/script/pprof.txtar b/cmd/cue/cmd/testdata/script/pprof.txtar
index 5056c01fe82..26516b5f083 100644
--- a/cmd/cue/cmd/testdata/script/pprof.txtar
+++ b/cmd/cue/cmd/testdata/script/pprof.txtar
@@ -15,10 +15,8 @@ exists cpu.out
 exists mem.out
 
 # The flags should still work with 'cue cmd', which runs a command within a command.
-# TODO: this is currently broken, as we try to start pprof twice.
 rm cpu.out
-! exec cue cmd --cpuprofile=cpu.out hello
-stderr 'cpu profiling already in use'
+exec cue cmd --cpuprofile=cpu.out hello
 go tool pprof -top cpu.out
 stdout 'Type: cpu'