Refactor BufferedLogger flushing and container command execution

- Add Flush(ctx, err) to BufferedLogger to automatically flush at
  ERROR on failure or DEBUG on success.
- Add executeWithLogger and executeWithOutput helpers to eliminate
  repetitive flush boilerplate.
- Unify runtime dependency log outputs into single log entries.
- Add unit tests and docstrings for BufferedLogger.

Author: shunping

sdks/go/container/tools/buffered_logging.go

@@ -52,9 +52,10 @@ func NewBufferedLoggerWithFlushInterval(ctx context.Context, logger *Logger, int
 	return &BufferedLogger{logger: logger, lastFlush: time.Now(), flushInterval: interval, periodicFlushContext: ctx, now: time.Now}
 }
 
-// Write implements the io.Writer interface, converting input to a string
-// and storing it in the BufferedLogger's buffer. If a logger is not provided,
-// the output is sent directly to os.Stderr.
+// Write implements the io.Writer interface. It buffers byte streams line-by-line
+// into memory and flushes periodically or upon calling Flush(), FlushAtError(), or
+// FlushAtDebug(). It is used primarily to redirect stdout/stderr of subprocesses or
+// standard Go log output. If a logger is not provided, the output is sent directly to os.Stderr.
 func (b *BufferedLogger) Write(p []byte) (int, error) {
 	if b.logger == nil {
 		return os.Stderr.Write(p)
@@ -86,6 +87,18 @@ func (b *BufferedLogger) Write(p []byte) (int, error) {
 	return n, err
 }
 
+// Flush flushes the contents of the buffer to the logging service.
+// If err is non-nil, it flushes at Error severity; otherwise it flushes at Debug severity.
+// It returns the provided error.
+func (b *BufferedLogger) Flush(ctx context.Context, err error) error {
+	if err != nil {
+		b.FlushAtError(ctx)
+	} else {
+		b.FlushAtDebug(ctx)
+	}
+	return err
+}
+
 // FlushAtError flushes the contents of the buffer to the logging
 // service at Error.
 func (b *BufferedLogger) FlushAtError(ctx context.Context) {
@@ -120,8 +133,9 @@ func (b *BufferedLogger) FlushAtDebug(ctx context.Context) {
 	b.lastFlush = time.Now()
 }
 
-// Prints directly to the logging service. If the logger is nil, prints directly to the
-// console. Used for the container pre-build workflow.
+// Printf directly writes formatted messages to the underlying logger/service,
+// bypassing line buffering. If the logger is nil, it prints directly to the
+// console. Used for direct informational logs and the container pre-build workflow.
 func (b *BufferedLogger) Printf(ctx context.Context, format string, args ...any) {
 	if b.logger == nil {
 		log.Printf(format, args...)

sdks/go/container/tools/buffered_logging_test.go

@@ -17,6 +17,7 @@ package tools
 
 import (
 	"context"
+	"errors"
 	"testing"
 	"time"
 
@@ -34,6 +35,24 @@ func getAllLogEntries(catcher *logCatcher) []*fnpb.LogEntry {
 func TestBufferedLogger(t *testing.T) {
 	ctx := context.Background()
 
+	t.Run("printf", func(t *testing.T) {
+		catcher := &logCatcher{}
+		l := &Logger{client: catcher}
+		bl := NewBufferedLogger(l)
+
+		bl.Printf(ctx, "test message")
+
+		received := catcher.msgs[0].GetLogEntries()[0]
+
+		if got, want := received.Message, "test message"; got != want {
+			t.Errorf("got message %q, want %q", got, want)
+		}
+
+		if got, want := received.Severity, fnpb.LogEntry_Severity_DEBUG; got != want {
+			t.Errorf("got severity %v, want %v", got, want)
+		}
+	})
+
 	t.Run("write", func(t *testing.T) {
 		catcher := &logCatcher{}
 		l := &Logger{client: catcher}
@@ -186,6 +205,55 @@ func TestBufferedLogger(t *testing.T) {
 		}
 	})
 
+	t.Run("flush with nil error", func(t *testing.T) {
+		catcher := &logCatcher{}
+		l := &Logger{client: catcher}
+		bl := NewBufferedLogger(l)
+
+		message := []byte("success message\n")
+		_, err := bl.Write(message)
+		if err != nil {
+			t.Fatalf("unexpected write error: %v", err)
+		}
+
+		if gotErr := bl.Flush(ctx, nil); gotErr != nil {
+			t.Errorf("Flush(ctx, nil) returned error %v, want nil", gotErr)
+		}
+
+		received := catcher.msgs[0].GetLogEntries()[0]
+		if got, want := received.Message, "success message"; got != want {
+			t.Errorf("got message %q, want %q", got, want)
+		}
+		if got, want := received.Severity, fnpb.LogEntry_Severity_DEBUG; got != want {
+			t.Errorf("got severity %v, want %v", got, want)
+		}
+	})
+
+	t.Run("flush with non-nil error", func(t *testing.T) {
+		catcher := &logCatcher{}
+		l := &Logger{client: catcher}
+		bl := NewBufferedLogger(l)
+
+		message := []byte("error message\n")
+		_, err := bl.Write(message)
+		if err != nil {
+			t.Fatalf("unexpected write error: %v", err)
+		}
+
+		originalErr := errors.New("command failed")
+		if gotErr := bl.Flush(ctx, originalErr); gotErr != originalErr {
+			t.Errorf("Flush(ctx, err) returned %v, want %v", gotErr, originalErr)
+		}
+
+		received := catcher.msgs[0].GetLogEntries()[0]
+		if got, want := received.Message, "error message"; got != want {
+			t.Errorf("got message %q, want %q", got, want)
+		}
+		if got, want := received.Severity, fnpb.LogEntry_Severity_ERROR; got != want {
+			t.Errorf("got severity %v, want %v", got, want)
+		}
+	})
+
 	t.Run("direct print", func(t *testing.T) {
 		catcher := &logCatcher{}
 		l := &Logger{client: catcher}

sdks/python/container/boot.go

@@ -18,7 +18,6 @@
 package main
 
 import (
-	"bytes"
 	"context"
 	"encoding/json"
 	"errors"
@@ -572,37 +571,27 @@ func logRuntimeDependencies(ctx context.Context, bufLogger *tools.BufferedLogger
 	if err != nil {
 		return err
 	}
-	bufLogger.Printf(ctx, "Python version in %s:", phase)
-	args := []string{"--version"}
-	if err := execx.ExecuteEnvWithIO(nil, os.Stdin, bufLogger, bufLogger, pythonVersion, args...); err != nil {
-		bufLogger.FlushAtError(ctx)
-	} else {
-		bufLogger.FlushAtDebug(ctx)
+	if out, err := executeWithOutput(ctx, bufLogger, pythonVersion, "--version"); err == nil {
+		bufLogger.Printf(ctx, "Python version in %s: %s", phase, strings.TrimSpace(string(out)))
 	}
-	bufLogger.Printf(ctx, "Dependencies in %s:", phase)
-	args = []string{"-m", "pip", "freeze", "--all"}
 
-	var stdout bytes.Buffer
-	if err := execx.ExecuteEnvWithIO(nil, os.Stdin, &stdout, bufLogger, pythonVersion, args...); err != nil {
-		bufLogger.FlushAtError(ctx)
-	} else {
-		bufLogger.FlushAtDebug(ctx)
-		bufLogger.Printf(ctx, "%s", stdout.String())
+	args := []string{"-m", "pip", "freeze", "--all"}
+	if out, err := executeWithOutput(ctx, bufLogger, pythonVersion, args...); err == nil {
+		bufLogger.Printf(ctx, "Dependencies in %s:\n%s", phase, string(out))
 	}
 	return nil
 }
 
 // logSubmissionEnvDependencies logs the python dependencies
 // installed in the submission environment.
 func logSubmissionEnvDependencies(ctx context.Context, bufLogger *tools.BufferedLogger, dir string) error {
-	bufLogger.Printf(ctx, "Dependencies in submission environment:")
 	// path for submission environment dependencies should match with the
 	// one defined in apache_beam/runners/portability/stager.py.
 	filename := filepath.Join(dir, "submission_environment_dependencies.txt")
 	content, err := os.ReadFile(filename)
 	if err != nil {
 		return err
 	}
-	bufLogger.Printf(ctx, "%s", string(content))
+	bufLogger.Printf(ctx, "Dependencies in submission environment:\n%s", string(content))
 	return nil
 }

sdks/python/container/piputil.go

@@ -41,6 +41,24 @@ var (
 const pipLogFlushInterval time.Duration = 15 * time.Second
 const unrecoverableURL string = "https://beam.apache.org/documentation/sdks/python-unrecoverable-errors/index.html#pip-dependency-resolution-failures"
 
+// executeWithLogger runs the program with os.Stdin, piping stdout and stderr to bufLogger,
+// and flushes bufLogger based on the execution result.
+func executeWithLogger(ctx context.Context, bufLogger *tools.BufferedLogger, prog string, args ...string) error {
+	err := execx.ExecuteEnvWithIO(nil, os.Stdin, bufLogger, bufLogger, prog, args...)
+	return bufLogger.Flush(ctx, err)
+}
+
+// executeWithOutput runs the program with os.Stdin, capturing stdout in a byte buffer
+// while piping stderr to bufLogger, and flushes bufLogger based on the execution result.
+func executeWithOutput(ctx context.Context, bufLogger *tools.BufferedLogger, prog string, args ...string) ([]byte, error) {
+	var stdout bytes.Buffer
+	err := execx.ExecuteEnvWithIO(nil, os.Stdin, &stdout, bufLogger, prog, args...)
+	if flushErr := bufLogger.Flush(ctx, err); flushErr != nil {
+		return nil, flushErr
+	}
+	return stdout.Bytes(), nil
+}
+
 // pipInstallRequirements installs the given requirement, if present.
 func pipInstallRequirements(ctx context.Context, logger *tools.Logger, files []string, dir, name string) error {
 	pythonVersion, err := expansionx.GetPythonVersion()
@@ -62,12 +80,9 @@ func pipInstallRequirements(ctx context.Context, logger *tools.Logger, files []s
 			// also installs dependencies. The key is that if all the packages have
 			// been installed in the first round then this command will be a no-op.
 			args = []string{"-m", "pip", "install", "-r", filepath.Join(dir, name), "--no-cache-dir", "--disable-pip-version-check", "--find-links", dir}
-			err := execx.ExecuteEnvWithIO(nil, os.Stdin, bufLogger, bufLogger, pythonVersion, args...)
-			if err != nil {
-				bufLogger.FlushAtError(ctx)
+			if err := executeWithLogger(ctx, bufLogger, pythonVersion, args...); err != nil {
 				return fmt.Errorf("PIP failed to install dependencies, got %s. This error may be unrecoverable, see %s for more information", err, unrecoverableURL)
 			}
-			bufLogger.FlushAtDebug(ctx)
 			return nil
 		}
 	}
@@ -121,23 +136,16 @@ func pipInstallPackage(ctx context.Context, logger *tools.Logger, files []string
 				if pipNoBuildIsolation {
 					args = append(args, "--no-build-isolation")
 				}
-				err := execx.ExecuteEnvWithIO(nil, os.Stdin, bufLogger, bufLogger, pythonVersion, args...)
-				if err != nil {
-					bufLogger.FlushAtError(ctx)
+				if err := executeWithLogger(ctx, bufLogger, pythonVersion, args...); err != nil {
 					return fmt.Errorf("PIP failed to install dependencies, got %s. This error may be unrecoverable, see %s for more information", err, unrecoverableURL)
-				} else {
-					bufLogger.FlushAtDebug(ctx)
 				}
 				args = []string{"-m", "pip", "install", "--no-cache-dir", "--disable-pip-version-check", filepath.Join(dir, packageSpec)}
 				if pipNoBuildIsolation {
 					args = append(args, "--no-build-isolation")
 				}
-				err = execx.ExecuteEnvWithIO(nil, os.Stdin, bufLogger, bufLogger, pythonVersion, args...)
-				if err != nil {
-					bufLogger.FlushAtError(ctx)
+				if err := executeWithLogger(ctx, bufLogger, pythonVersion, args...); err != nil {
 					return fmt.Errorf("PIP failed to install dependencies, got %s. This error may be unrecoverable, see %s for more information", err, unrecoverableURL)
 				}
-				bufLogger.FlushAtDebug(ctx)
 				return nil
 			}
 
@@ -146,12 +154,9 @@ func pipInstallPackage(ctx context.Context, logger *tools.Logger, files []string
 			if pipNoBuildIsolation {
 				args = append(args, "--no-build-isolation")
 			}
-			err := execx.ExecuteEnvWithIO(nil, os.Stdin, bufLogger, bufLogger, pythonVersion, args...)
-			if err != nil {
-				bufLogger.FlushAtError(ctx)
+			if err := executeWithLogger(ctx, bufLogger, pythonVersion, args...); err != nil {
 				return fmt.Errorf("PIP failed to install dependencies, got %s. This error may be unrecoverable, see %s for more information", err, unrecoverableURL)
 			}
-			bufLogger.FlushAtDebug(ctx)
 			return nil
 		}
 	}