From cde1b520e213e8e7aad36f1264d20311fa014392 Mon Sep 17 00:00:00 2001
From: Ignas Anikevicius <240938+aignas@users.noreply.github.com>
Date: Fri, 17 Nov 2023 12:39:03 +0900
Subject: [PATCH] fix(gazelle): make cmd.Wait more idiomatic (#1550)

It seems that the documentation for the `cmd.Wait` explicitly
asks the users to not wait on the command immediately after
starting because it may close pipes too early and cause
unintended side-effects as described in #1546.

Fixes #1546.

Co-authored-by: Richard Levasseur <rlevasseur@google.com>
---
 CHANGELOG.md                  |  3 +++
 gazelle/python/parser.go      | 22 ++++++++++------------
 gazelle/python/std_modules.go | 24 +++++++++++-------------
 3 files changed, 24 insertions(+), 25 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index d6b419d08c..32ab939c5d 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -97,6 +97,9 @@ Breaking changes:
 * (gazelle) Generate a single `py_test` target when `gazelle:python_generation_mode project`
   is used.
 
+* (gazelle) Move waiting for the Python interpreter process to exit to the shutdown hook
+  to make the usage of the `exec.Command` more idiomatic.
+
 * (toolchains) Keep tcl subdirectory in Windows build of hermetic interpreter.
 
 * (bzlmod) sub-modules now don't have the `//conditions:default` clause in the
diff --git a/gazelle/python/parser.go b/gazelle/python/parser.go
index ad55e03a01..89310267c3 100644
--- a/gazelle/python/parser.go
+++ b/gazelle/python/parser.go
@@ -32,6 +32,7 @@ import (
 )
 
 var (
+	parserCmd    *exec.Cmd
 	parserStdin  io.WriteCloser
 	parserStdout io.Reader
 	parserMutex  sync.Mutex
@@ -40,40 +41,37 @@ var (
 func startParserProcess(ctx context.Context) {
 	// due to #691, we need a system interpreter to boostrap, part of which is
 	// to locate the hermetic interpreter.
-	cmd := exec.CommandContext(ctx, "python3", helperPath, "parse")
-	cmd.Stderr = os.Stderr
+	parserCmd = exec.CommandContext(ctx, "python3", helperPath, "parse")
+	parserCmd.Stderr = os.Stderr
 
-	stdin, err := cmd.StdinPipe()
+	stdin, err := parserCmd.StdinPipe()
 	if err != nil {
 		log.Printf("failed to initialize parser: %v\n", err)
 		os.Exit(1)
 	}
 	parserStdin = stdin
 
-	stdout, err := cmd.StdoutPipe()
+	stdout, err := parserCmd.StdoutPipe()
 	if err != nil {
 		log.Printf("failed to initialize parser: %v\n", err)
 		os.Exit(1)
 	}
 	parserStdout = stdout
 
-	if err := cmd.Start(); err != nil {
+	if err := parserCmd.Start(); err != nil {
 		log.Printf("failed to initialize parser: %v\n", err)
 		os.Exit(1)
 	}
-
-	go func() {
-		if err := cmd.Wait(); err != nil {
-			log.Printf("failed to wait for parser: %v\n", err)
-			os.Exit(1)
-		}
-	}()
 }
 
 func shutdownParserProcess() {
 	if err := parserStdin.Close(); err != nil {
 		fmt.Fprintf(os.Stderr, "error closing parser: %v", err)
 	}
+
+	if err := parserCmd.Wait(); err != nil {
+		log.Printf("failed to wait for parser: %v\n", err)
+	}
 }
 
 // python3Parser implements a parser for Python files that extracts the modules
diff --git a/gazelle/python/std_modules.go b/gazelle/python/std_modules.go
index dd59cd8832..8a016afed6 100644
--- a/gazelle/python/std_modules.go
+++ b/gazelle/python/std_modules.go
@@ -29,6 +29,7 @@ import (
 )
 
 var (
+	stdModulesCmd    *exec.Cmd
 	stdModulesStdin  io.WriteCloser
 	stdModulesStdout io.Reader
 	stdModulesMutex  sync.Mutex
@@ -40,42 +41,39 @@ func startStdModuleProcess(ctx context.Context) {
 
 	// due to #691, we need a system interpreter to boostrap, part of which is
 	// to locate the hermetic interpreter.
-	cmd := exec.CommandContext(ctx, "python3", helperPath, "std_modules")
-	cmd.Stderr = os.Stderr
+	stdModulesCmd = exec.CommandContext(ctx, "python3", helperPath, "std_modules")
+	stdModulesCmd.Stderr = os.Stderr
 	// All userland site-packages should be ignored.
-	cmd.Env = []string{"PYTHONNOUSERSITE=1"}
+	stdModulesCmd.Env = []string{"PYTHONNOUSERSITE=1"}
 
-	stdin, err := cmd.StdinPipe()
+	stdin, err := stdModulesCmd.StdinPipe()
 	if err != nil {
 		log.Printf("failed to initialize std_modules: %v\n", err)
 		os.Exit(1)
 	}
 	stdModulesStdin = stdin
 
-	stdout, err := cmd.StdoutPipe()
+	stdout, err := stdModulesCmd.StdoutPipe()
 	if err != nil {
 		log.Printf("failed to initialize std_modules: %v\n", err)
 		os.Exit(1)
 	}
 	stdModulesStdout = stdout
 
-	if err := cmd.Start(); err != nil {
+	if err := stdModulesCmd.Start(); err != nil {
 		log.Printf("failed to initialize std_modules: %v\n", err)
 		os.Exit(1)
 	}
-
-	go func() {
-		if err := cmd.Wait(); err != nil {
-			log.Printf("failed to wait for std_modules: %v\n", err)
-			os.Exit(1)
-		}
-	}()
 }
 
 func shutdownStdModuleProcess() {
 	if err := stdModulesStdin.Close(); err != nil {
 		fmt.Fprintf(os.Stderr, "error closing std module: %v", err)
 	}
+
+	if err := stdModulesCmd.Wait(); err != nil {
+		log.Printf("failed to wait for std_modules: %v\n", err)
+	}
 }
 
 func isStdModule(m module) (bool, error) {