compiler_wrapper: refactor rusage_flag

Refactors rusage_flag and its usage in compiler_wrapper to support
the use of TOOLCHAIN_RUSAGE_OUTPUT with either FORCE_DISABLE_WERROR.

BUG=chromium:1167958
TEST=Modified and ran unit tests

Change-Id: I9a643ceabbbeca3bf3a94270d4f4b4d814812f93
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/third_party/toolchain-utils/+/2635451
Tested-by: Ryan Beltran <ryanbeltran@chromium.org>
Reviewed-by: George Burgess <gbiv@chromium.org>

4 files changed, 162 insertions, 93 deletions

diff --git a/compiler_wrapper/compiler_wrapper.go b/compiler_wrapper/compiler_wrapper.go
index aca5aa7e..631f80fb 100644
--- a/compiler_wrapper/compiler_wrapper.go
+++ b/compiler_wrapper/compiler_wrapper.go
@@ -141,41 +141,53 @@ func callCompilerInternal(env env, cfg *config, inputCmd *command) (exitCode int
 			}
 		}
 	}
+
 	rusageLogfileName := getRusageLogFilename(env)
 	bisectStage := getBisectStage(env)
+
+	if rusageLogfileName != "" {
+		compilerCmd = removeRusageFromCommand(compilerCmd)
+	}
+
 	if shouldForceDisableWerror(env, cfg) {
-		if rusageLogfileName != "" {
-			return 0, newUserErrorf("TOOLCHAIN_RUSAGE_OUTPUT is meaningless with FORCE_DISABLE_WERROR")
-		}
 		if bisectStage != "" {
 			return 0, newUserErrorf("BISECT_STAGE is meaningless with FORCE_DISABLE_WERROR")
 		}
-		return doubleBuildWithWNoError(env, cfg, compilerCmd)
+		return doubleBuildWithWNoError(env, cfg, compilerCmd, rusageLogfileName)
 	}
 	if shouldCompileWithFallback(env) {
 		if rusageLogfileName != "" {
-			return 0, newUserErrorf("TOOLCHAIN_RUSAGE_OUTPUT is meaningless with FORCE_DISABLE_WERROR")
+			return 0, newUserErrorf("TOOLCHAIN_RUSAGE_OUTPUT is meaningless with ANDROID_LLVM_PREBUILT_COMPILER_PATH")
 		}
 		if bisectStage != "" {
-			return 0, newUserErrorf("BISECT_STAGE is meaningless with FORCE_DISABLE_WERROR")
+			return 0, newUserErrorf("BISECT_STAGE is meaningless with ANDROID_LLVM_PREBUILT_COMPILER_PATH")
 		}
 		return compileWithFallback(env, cfg, compilerCmd, mainBuilder.absWrapperPath)
 	}
-	if rusageLogfileName != "" {
-		if bisectStage != "" {
-			return 0, newUserErrorf("BISECT_STAGE is meaningless with TOOLCHAIN_RUSAGE_OUTPUT")
-		}
-		return logRusage(env, rusageLogfileName, compilerCmd)
-	}
 	if bisectStage != "" {
+		if rusageLogfileName != "" {
+			return 0, newUserErrorf("TOOLCHAIN_RUSAGE_OUTPUT is meaningless with BISECT_STAGE")
+		}
 		compilerCmd, err = calcBisectCommand(env, cfg, bisectStage, compilerCmd)
 		if err != nil {
 			return 0, err
 		}
 	}
-	// Note: We return an exit code only if the underlying env is not
-	// really doing an exec, e.g. commandRecordingEnv.
-	return wrapSubprocessErrorWithSourceLoc(compilerCmd, env.exec(compilerCmd))
+
+	commitRusage, err := maybeCaptureRusage(env, rusageLogfileName, compilerCmd, func() error {
+		// Note: We return an exit code only if the underlying env is not
+		// really doing an exec, e.g. commandRecordingEnv.
+		exitCode, err = wrapSubprocessErrorWithSourceLoc(compilerCmd, env.exec(compilerCmd))
+		return err
+	})
+	if err != nil {
+		return exitCode, err
+	}
+	if err := commitRusage(exitCode); err != nil {
+		return exitCode, fmt.Errorf("commiting rusage: %v", err)
+	}
+
+	return exitCode, err
 }
 
 func prepareClangCommand(builder *commandBuilder) (err error) {
diff --git a/compiler_wrapper/compiler_wrapper_test.go b/compiler_wrapper/compiler_wrapper_test.go
index 3584f6ed..680312f9 100644
--- a/compiler_wrapper/compiler_wrapper_test.go
+++ b/compiler_wrapper/compiler_wrapper_test.go
@@ -9,6 +9,8 @@ import (
 	"errors"
 	"fmt"
 	"io"
+	"os"
+	"path/filepath"
 	"strings"
 	"syscall"
 	"testing"
@@ -94,15 +96,33 @@ func TestLogMissingCCacheExecError(t *testing.T) {
 	})
 }
 
-func TestErrorOnLogRusageAndForceDisableWError(t *testing.T) {
+func TestLogRusageAndForceDisableWError(t *testing.T) {
 	withTestContext(t, func(ctx *testContext) {
+		logFileName := filepath.Join(ctx.tempDir, "rusage.log")
 		ctx.env = []string{
 			"FORCE_DISABLE_WERROR=1",
-			"TOOLCHAIN_RUSAGE_OUTPUT=rusage.log",
+			"TOOLCHAIN_RUSAGE_OUTPUT=" + logFileName,
 		}
-		stderr := ctx.mustFail(callCompiler(ctx, ctx.cfg, ctx.newCommand(gccX86_64, mainCc)))
-		if err := verifyNonInternalError(stderr, "TOOLCHAIN_RUSAGE_OUTPUT is meaningless with FORCE_DISABLE_WERROR"); err != nil {
-			t.Error(err)
+		ctx.cmdMock = func(cmd *command, stdin io.Reader, stdout io.Writer, stderr io.Writer) error {
+			switch ctx.cmdCount {
+			case 1:
+				io.WriteString(stderr, "-Werror originalerror")
+				return newExitCodeError(1)
+			case 2:
+				return nil
+			default:
+				t.Fatalf("unexpected command: %#v", cmd)
+				return nil
+			}
+		}
+		ctx.must(callCompiler(ctx, ctx.cfg, ctx.newCommand(gccX86_64, mainCc)))
+		if _, err := os.Stat(logFileName); os.IsNotExist(err) {
+			t.Errorf("no logfile created at TOOLCHAIN_RUSAGE_OUTPUT path %q", logFileName)
+		} else if err != nil {
+			t.Fatalf("error checking for rusage logfile at %q: %v", logFileName, err)
+		}
+		if ctx.cmdCount != 2 {
+			t.Errorf("expected 2 calls. Got: %d", ctx.cmdCount)
 		}
 	})
 }
@@ -114,7 +134,7 @@ func TestErrorOnLogRusageAndBisect(t *testing.T) {
 			"TOOLCHAIN_RUSAGE_OUTPUT=rusage.log",
 		}
 		stderr := ctx.mustFail(callCompiler(ctx, ctx.cfg, ctx.newCommand(gccX86_64, mainCc)))
-		if err := verifyNonInternalError(stderr, "BISECT_STAGE is meaningless with TOOLCHAIN_RUSAGE_OUTPUT"); err != nil {
+		if err := verifyNonInternalError(stderr, "TOOLCHAIN_RUSAGE_OUTPUT is meaningless with BISECT_STAGE"); err != nil {
 			t.Error(err)
 		}
 	})
diff --git a/compiler_wrapper/disable_werror_flag.go b/compiler_wrapper/disable_werror_flag.go
index 8f20b6f3..eacf9530 100644
--- a/compiler_wrapper/disable_werror_flag.go
+++ b/compiler_wrapper/disable_werror_flag.go
@@ -56,7 +56,7 @@ func isLikelyAConfTest(cfg *config, cmd *command) bool {
 	return false
 }
 
-func doubleBuildWithWNoError(env env, cfg *config, originalCmd *command) (exitCode int, err error) {
+func doubleBuildWithWNoError(env env, cfg *config, originalCmd *command, rusageLogfileName string) (exitCode int, err error) {
 	originalStdoutBuffer := &bytes.Buffer{}
 	originalStderrBuffer := &bytes.Buffer{}
 	// TODO: This is a bug in the old wrapper that it drops the ccache path
@@ -70,8 +70,12 @@ func doubleBuildWithWNoError(env env, cfg *config, originalCmd *command) (exitCo
 		return 0, wrapErrorwithSourceLocf(err, "prebuffering stdin: %v", err)
 	}
 
-	originalExitCode, err := wrapSubprocessErrorWithSourceLoc(originalCmd,
-		env.run(originalCmd, getStdin(), originalStdoutBuffer, originalStderrBuffer))
+	var originalExitCode int
+	commitOriginalRusage, err := maybeCaptureRusage(env, rusageLogfileName, originalCmd, func() error {
+		originalExitCode, err = wrapSubprocessErrorWithSourceLoc(originalCmd,
+			env.run(originalCmd, getStdin(), originalStdoutBuffer, originalStderrBuffer))
+		return err
+	})
 	if err != nil {
 		return 0, err
 	}
@@ -85,6 +89,9 @@ func doubleBuildWithWNoError(env env, cfg *config, originalCmd *command) (exitCo
 			bytes.Contains(originalStdoutBufferBytes, []byte("warnings-as-errors")) ||
 			bytes.Contains(originalStdoutBufferBytes, []byte("clang-diagnostic-")))
 	if !shouldRetry {
+		if err := commitOriginalRusage(originalExitCode); err != nil {
+			return 0, fmt.Errorf("commiting rusage: %v", err)
+		}
 		originalStdoutBuffer.WriteTo(env.stdout())
 		originalStderrBuffer.WriteTo(env.stderr())
 		return originalExitCode, nil
@@ -97,19 +104,32 @@ func doubleBuildWithWNoError(env env, cfg *config, originalCmd *command) (exitCo
 		Args:       disableWerrorFlags(originalCmd.Args),
 		EnvUpdates: originalCmd.EnvUpdates,
 	}
-	retryExitCode, err := wrapSubprocessErrorWithSourceLoc(retryCommand,
-		env.run(retryCommand, getStdin(), retryStdoutBuffer, retryStderrBuffer))
+
+	var retryExitCode int
+	commitRetryRusage, err := maybeCaptureRusage(env, rusageLogfileName, retryCommand, func() error {
+		retryExitCode, err = wrapSubprocessErrorWithSourceLoc(retryCommand,
+			env.run(retryCommand, getStdin(), retryStdoutBuffer, retryStderrBuffer))
+		return err
+	})
 	if err != nil {
 		return 0, err
 	}
+
 	// If -Wno-error fixed us, pretend that we never ran without -Wno-error. Otherwise, pretend
 	// that we never ran the second invocation.
 	if retryExitCode != 0 {
 		originalStdoutBuffer.WriteTo(env.stdout())
 		originalStderrBuffer.WriteTo(env.stderr())
+		if err := commitOriginalRusage(originalExitCode); err != nil {
+			return 0, fmt.Errorf("commiting rusage: %v", err)
+		}
 		return originalExitCode, nil
 	}
 
+	if err := commitRetryRusage(retryExitCode); err != nil {
+		return 0, fmt.Errorf("commiting rusage: %v", err)
+	}
+
 	retryStdoutBuffer.WriteTo(env.stdout())
 	retryStderrBuffer.WriteTo(env.stderr())
 
diff --git a/compiler_wrapper/rusage_flag.go b/compiler_wrapper/rusage_flag.go
index c6f887d7..a2159021 100644
--- a/compiler_wrapper/rusage_flag.go
+++ b/compiler_wrapper/rusage_flag.go
@@ -49,89 +49,106 @@ type rusageLog struct {
 	WorkingDirectory string   `json:"working_directory"`
 }
 
-func logRusage(env env, logFileName string, compilerCmd *command) (exitCode int, err error) {
-	rusageBefore := syscall.Rusage{}
-	if err := syscall.Getrusage(syscall.RUSAGE_CHILDREN, &rusageBefore); err != nil {
-		return 0, err
-	}
-	compilerCmdWithoutRusage := &command{
+func removeRusageFromCommand(compilerCmd *command) *command {
+	return &command{
 		Path:       compilerCmd.Path,
 		Args:       compilerCmd.Args,
 		EnvUpdates: append(compilerCmd.EnvUpdates, "TOOLCHAIN_RUSAGE_OUTPUT="),
 	}
+}
+
+// maybeCaptureRusage captures rusage for execution of action()
+// 	unless action returns an error or logFileName is ""
+// a function is returned which saves the rusage log data at logFileName unless logFileName is ""
+// an error is returned if action returns an error, or rusage commands in syscall fails
+func maybeCaptureRusage(env env, logFileName string, compilerCmd *command, action func() error) (maybeCommitToFile func(exitCode int) error, err error) {
+	if logFileName == "" {
+		if err := action(); err != nil {
+			return nil, err
+		}
+		return func(int) error { return nil }, nil
+	}
+
+	rusageBefore := syscall.Rusage{}
+	if err := syscall.Getrusage(syscall.RUSAGE_CHILDREN, &rusageBefore); err != nil {
+		return nil, fmt.Errorf("getting initial rusage: %v", err)
+	}
 	startTime := time.Now()
-	exitCode, err = wrapSubprocessErrorWithSourceLoc(compilerCmdWithoutRusage,
-		env.run(compilerCmdWithoutRusage, env.stdin(), env.stdout(), env.stderr()))
-	if err != nil {
-		return 0, err
+
+	if err := action(); err != nil {
+		return nil, err
 	}
+
 	elapsedRealTime := time.Since(startTime)
 	rusageAfter := syscall.Rusage{}
 	if err := syscall.Getrusage(syscall.RUSAGE_CHILDREN, &rusageAfter); err != nil {
-		return 0, err
-	}
-	elapsedSysTime := time.Duration(rusageAfter.Stime.Nano()-rusageBefore.Stime.Nano()) * time.Nanosecond
-	elapsedUserTime := time.Duration(rusageAfter.Utime.Nano()-rusageBefore.Utime.Nano()) * time.Nanosecond
-	// Note: We assume that the compiler takes more heap than any other
-	// subcommands that we might have executed before.
-	maxMemUsed := rusageAfter.Maxrss
-	absCompilerPath := getAbsCmdPath(env, compilerCmd)
-
-	// We need to temporarily set umask to 0 to ensure 777 permissions are actually 777
-	// This effects builderbots in particular
-	oldMask := syscall.Umask(0)
-	defer syscall.Umask(oldMask)
-
-	if err := os.MkdirAll(filepath.Dir(logFileName), 0777); err != nil {
-		return 0, wrapErrorwithSourceLocf(err, "error creating rusage log directory %s", logFileName)
+		return nil, fmt.Errorf("getting final rusage: %v", err)
 	}
 
-	timeUnit := float64(time.Second)
+	return func(exitCode int) error {
+		elapsedSysTime := time.Duration(rusageAfter.Stime.Nano()-rusageBefore.Stime.Nano()) * time.Nanosecond
+		elapsedUserTime := time.Duration(rusageAfter.Utime.Nano()-rusageBefore.Utime.Nano()) * time.Nanosecond
+		// Note: We assume that the compiler takes more heap than any other
+		// subcommands that we might have executed before.
+		maxMemUsed := rusageAfter.Maxrss
+		absCompilerPath := getAbsCmdPath(env, compilerCmd)
+
+		// We need to temporarily set umask to 0 to ensure 777 permissions are actually 777
+		// This effects builderbots in particular
+		oldMask := syscall.Umask(0)
+		defer syscall.Umask(oldMask)
+
+		// We want to know what package is being compiled. The working directory gives us a good clue.
+		cwd, err := os.Getwd()
+		if err != nil {
+			return wrapErrorwithSourceLocf(err, "error getting working directory for rusage log")
+		}
 
-	// We want to know what package is being compiled. The working directory gives us a good clue.
-	cwd, err := os.Getwd()
-	if err != nil {
-		return 0, wrapErrorwithSourceLocf(err, "error getting working directory for rusage log")
-	}
+		if err := os.MkdirAll(filepath.Dir(logFileName), 0777); err != nil {
+			return wrapErrorwithSourceLocf(err, "error creating rusage log directory %s", logFileName)
+		}
 
-	logEntry := rusageLog{
-		ExitCode:         exitCode,
-		ElapsedRealTime:  float64(elapsedRealTime) / timeUnit,
-		ElapsedUserTime:  float64(elapsedUserTime) / timeUnit,
-		ElapsedSysTime:   float64(elapsedSysTime) / timeUnit,
-		MaxMemUsed:       maxMemUsed,
-		Compiler:         absCompilerPath,
-		CompilerArgs:     compilerCmd.Args,
-		WorkingDirectory: cwd,
-	}
+		timeUnit := float64(time.Second)
+
+		logEntry := rusageLog{
+			ExitCode:         exitCode,
+			ElapsedRealTime:  float64(elapsedRealTime) / timeUnit,
+			ElapsedUserTime:  float64(elapsedUserTime) / timeUnit,
+			ElapsedSysTime:   float64(elapsedSysTime) / timeUnit,
+			MaxMemUsed:       maxMemUsed,
+			Compiler:         absCompilerPath,
+			CompilerArgs:     compilerCmd.Args,
+			WorkingDirectory: cwd,
+		}
 
-	// Note: using file mode 0666 so that a root-created log is writable by others.
-	logFile, err := os.OpenFile(logFileName, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0666)
-	if err != nil {
-		return 0, wrapErrorwithSourceLocf(err, "creating rusage logfile %s", logFileName)
-	}
+		// Note: using file mode 0666 so that a root-created log is writable by others.
+		logFile, err := os.OpenFile(logFileName, os.O_CREATE|os.O_APPEND|os.O_WRONLY, 0666)
+		if err != nil {
+			return wrapErrorwithSourceLocf(err, "creating rusage logfile %s", logFileName)
+		}
 
-	// O_APPEND's atomicity guarantees are only for writes up to a certain size. If we don't
-	// lock the file, we might end up with corrupted records.
-	//
-	// Note that Close()'ing the file releases all associated locks.
-	if err := lockFileExclusive(logFile.Fd()); err != nil {
-		_ = logFile.Close()
-		return 0, wrapErrorwithSourceLocf(err, "locking rusage logfile %s: %v", logFileName, err)
-	}
+		// O_APPEND's atomicity guarantees are only for writes up to a certain size. If we don't
+		// lock the file, we might end up with corrupted records.
+		//
+		// Note that Close()'ing the file releases all associated locks.
+		if err := lockFileExclusive(logFile.Fd()); err != nil {
+			_ = logFile.Close()
+			return wrapErrorwithSourceLocf(err, "locking rusage logfile %s: %v", logFileName, err)
+		}
 
-	if err := json.NewEncoder(logFile).Encode(logEntry); err != nil {
-		_ = logFile.Close()
-		return 0, wrapErrorwithSourceLocf(err, "converting rusage logfile entry to JSON %v", logEntry)
-	}
+		if err := json.NewEncoder(logFile).Encode(logEntry); err != nil {
+			_ = logFile.Close()
+			return wrapErrorwithSourceLocf(err, "converting rusage logfile entry to JSON %v", logEntry)
+		}
 
-	closeErr := logFile.Close()
-	if err != nil {
-		return 0, wrapErrorwithSourceLocf(err, "writing to rusage logfile %s: %v", logFileName, err)
-	}
-	if closeErr != nil {
-		return 0, wrapErrorwithSourceLocf(err, "closing rusage logfile %s: %v", logFileName, closeErr)
-	}
+		closeErr := logFile.Close()
+		if err != nil {
+			return wrapErrorwithSourceLocf(err, "writing to rusage logfile %s: %v", logFileName, err)
+		}
+		if closeErr != nil {
+			return wrapErrorwithSourceLocf(err, "closing rusage logfile %s: %v", logFileName, closeErr)
+		}
 
-	return exitCode, nil
+		return nil
+	}, nil
 }