fix: handle send timeouts in bes_pipe and bes_backend

gregmagolan · gregmagolan · commit bb4cb2a143bf · 2025-11-19T22:17:26.000-08:00
diff --git a/pkg/plugin/system/bep/bes_backend.go b/pkg/plugin/system/bep/bes_backend.go
@@ -26,6 +26,7 @@ import (
 	"os"
 	"strings"
 	"sync"
+	"time"
 
 	"github.com/fatih/color"
 	"github.com/golang/protobuf/ptypes/empty"
@@ -176,7 +177,7 @@ func (bb *besBackend) GracefulStop() {
 }
 
 // Addr returns the address for the gRPC server. Since the address is determined
-// by the OS based on an available port at the time the gRPC server starts, this
+// by the OS based on available port at the time the gRPC server starts, this
 // method returns the address to be used to construct the `bes_backend` flag
 // passed to the `bazel (build|test|run)` commands. The address includes the
 // scheme (protocol).
@@ -232,7 +233,7 @@ func (bb *besBackend) PublishLifecycleEvent(
 	broadCastEvent := func(ctx context.Context, req *buildv1.PublishLifecycleEventRequest) {
 		eg, ctx := errgroup.WithContext(ctx)
 		for _, p := range bb.besProxies {
-			if !p.Healthy() {
+			if !p.IsHealthy() {
 				continue
 			}
 
@@ -297,7 +298,7 @@ func (bb *besBackend) SendEventsToSubscribers(c <-chan *buildv1.PublishBuildTool
 	}
 }
 
-func (bb *besBackend) setupBesUpstreamBackends(ctx context.Context, optionsparsed *buildeventstream.OptionsParsed) error {
+func (bb *besBackend) setupBesUpstreamBackends(ctx context.Context, optionsParsed *buildeventstream.OptionsParsed) error {
 	defer close(bb.ready)
 
 	backends := []string{}
@@ -308,7 +309,7 @@ func (bb *besBackend) setupBesUpstreamBackends(ctx context.Context, optionsparse
 	hasNoWaitForUpload := false
 
 	// Parse backends first to build up that map, then parse headers
-	for _, arg := range optionsparsed.CmdLine {
+	for _, arg := range optionsParsed.CmdLine {
 		if strings.HasPrefix(arg, "--bes_backend=") {
 			// Always skip our bes_backend to avoid recursive uploads.
 			if arg == fmt.Sprintf("--bes_backend=%s", bb.Addr()) {
@@ -448,7 +449,7 @@ func (bb *besBackend) PublishBuildToolEventStream(
 					switch event.Id.Id.(type) {
 					case *buildeventstream.BuildEventId_OptionsParsed:
 						// Received options event, setup bes upstream backends based off commandline arguments bazel reported.
-						// setup upstream backends async to prevent bazel client from waiting for upstream bes connections.
+						// setup upstream backends async to prevent bazel to waiting for upstream bes connections.
 						eg.Go(func() error {
 							err = bb.setupBesUpstreamBackends(egCtx, event.GetOptionsParsed())
 							if err != nil {
@@ -496,14 +497,14 @@ func (bb *besBackend) PublishBuildToolEventStream(
 		<-bb.ready
 		// Goroutines to receive acks from BES proxies
 		for _, bp := range bb.besProxies {
-			if !bp.Healthy() {
+			if !bp.IsHealthy() {
 				continue
 			}
 			proxy := bp // make a copy of the BESProxy before passing into the go-routine below.
 			eg.Go(func() error {
 				for {
 					// If the proxy is not healthy, break out of the loop
-					if !proxy.Healthy() {
+					if !proxy.IsHealthy() {
 						break
 					}
 					_, err := proxy.Recv()
@@ -527,19 +528,46 @@ func (bb *besBackend) PublishBuildToolEventStream(
 	// Goroutine to forward to build event to BES proxies
 	eg.Go(func() error {
 		for fwd := range fwdChanRead {
+			egFwd := errgroup.Group{}
+
 			for _, bp := range bb.besProxies {
-				if !bp.Healthy() {
-					continue
-				}
-				err := bp.Send(fwd)
-				if err != nil {
-					// If we fail to send to a proxy then print out an error but don't fail the GRPC call
-					fmt.Fprintf(os.Stderr, "Error sending build event to %v: %s\n", bp.Host(), err.Error())
-				}
+				bp := bp // capture
+				egFwd.Go(func() error {
+					if !bp.IsHealthy() {
+						return nil
+					}
+
+					// Channel for Send result
+					sendCh := make(chan error, 1)
+
+					// Run Send in goroutine
+					go func() {
+						err := bp.Send(fwd)
+						sendCh <- err
+					}()
+
+					// Wait for Send or timeout
+					select {
+					case err := <-sendCh:
+						if err != nil {
+							fmt.Fprintf(os.Stderr, "Error sending build event to %v: %s\n", bp.Host(), err.Error())
+							bp.MarkUnhealthy()
+						}
+						return nil
+					case <-time.After(sendTimeout):
+						fmt.Fprintf(os.Stderr, "Timeout sending build event to %v: marking unhealthy\n", bp.Host())
+						bp.MarkUnhealthy()
+						return nil
+					}
+				})
+			}
+
+			if err := egFwd.Wait(); err != nil {
+				// Optionally handle errors from sends, but since we log inside, perhaps no need to propagate
 			}
 		}
 		for _, bp := range bb.besProxies {
-			if !bp.Healthy() {
+			if !bp.IsHealthy() {
 				continue
 			}
 			if err := bp.CloseSend(); err != nil {
diff --git a/pkg/plugin/system/bep/bes_backend_test.go b/pkg/plugin/system/bep/bes_backend_test.go
@@ -500,7 +500,7 @@ func TestPublishBuildToolEventStream(t *testing.T) {
 
 		besProxy.
 			EXPECT().
-			Healthy().
+			IsHealthy().
 			Return(true).
 			Times(5)
 		besProxy.
diff --git a/pkg/plugin/system/bep/bes_pipe.go b/pkg/plugin/system/bep/bes_pipe.go
@@ -49,6 +49,7 @@ type BESPipeInterceptor interface {
 const besEventGlobalTimeoutDuration = 5 * time.Minute
 const besEventThrottleDuration = 50 * time.Millisecond
 const gracePeriodDuration = 2 * time.Second
+const sendTimeout = 10 * time.Second // Timeout for individual Send calls to detect backed-up receivers
 
 func NewBESPipe(buildId, invocationId string) (BESPipeInterceptor, error) {
 	return &besPipe{
@@ -72,6 +73,9 @@ type besPipe struct {
 	besInvocationId string
 	besProxies      []besproxy.BESProxy
 
+	// Track whether we have already unlinked the pipe due to backend failure
+	pipeAborted sync.Once
+
 	wg *sync.WaitGroup
 }
 
@@ -101,7 +105,7 @@ func (bb *besPipe) RegisterBesProxy(ctx context.Context, p besproxy.BESProxy) {
 	go func() {
 		for {
 			// If the proxy is not healthy, break out of the loop
-			if !p.Healthy() {
+			if !p.IsHealthy() {
 				break
 			}
 			_, err := p.Recv()
@@ -116,6 +120,8 @@ func (bb *besPipe) RegisterBesProxy(ctx context.Context, p besproxy.BESProxy) {
 				break
 			}
 		}
+		// When the ACK goroutine exits (usually because of error), check if we should abort the pipe
+		bb.maybeAbortPipeBecauseNoHealthyBackends()
 	}()
 }
 
@@ -156,32 +162,35 @@ func (bb *besPipe) ServeWait(ctx context.Context) error {
 	bb.wg.Add(1)
 	go func() {
 		defer bb.wg.Done()
+
+		// If the overall context is cancelled, abort the pipe immediately
+		go func() {
+			<-ctx.Done()
+			bb.maybeAbortPipeBecauseNoHealthyBackends()
+		}()
+
 		conn, err := os.OpenFile(bb.bepBinPath, os.O_RDONLY, os.ModeNamedPipe)
 		if err != nil {
 			bb.errorsMutex.Lock()
-			defer bb.errorsMutex.Unlock()
 			bb.errors.Insert(fmt.Errorf("failed to accept connection on BES pipe %s: %w", bb.bepBinPath, err))
+			bb.errorsMutex.Unlock()
 			return
 		}
-
-		defer func() {
-			conn.Close()
-		}()
+		defer conn.Close()
 
 		if err := bb.streamBesEvents(ctx, conn); err != nil {
 			bb.errorsMutex.Lock()
-			defer bb.errorsMutex.Unlock()
 			bb.errors.Insert(fmt.Errorf("failed to stream BES events: %w", err))
+			bb.errorsMutex.Unlock()
 			return
 		}
 
+		// Normal completion path
 		for _, p := range bb.besProxies {
-			if !p.Healthy() {
+			if !p.IsHealthy() {
 				continue
 			}
-
 			bb.sendFinalLifecycleEvents(context.Background(), p)
-
 			if err := p.CloseSend(); err != nil {
 				fmt.Fprintf(os.Stderr, "Error closing build event stream to %v: %s\n", p.Host(), err.Error())
 			}
@@ -190,25 +199,47 @@ func (bb *besPipe) ServeWait(ctx context.Context) error {
 	return nil
 }
 
+// maybeAbortPipeBecauseNoHealthyBackends unlinks the FIFO if all backends are unhealthy.
+// This gives Bazel a broken pipe → it aborts the upload and exits.
+func (bb *besPipe) maybeAbortPipeBecauseNoHealthyBackends() {
+	if len(bb.besProxies) == 0 {
+		return
+	}
+
+	var anyHealthy bool
+	for _, p := range bb.besProxies {
+		if p.IsHealthy() {
+			anyHealthy = true
+			break
+		}
+	}
+	if anyHealthy {
+		return
+	}
+
+	bb.pipeAborted.Do(func() {
+		fmt.Fprintf(os.Stderr, "All BES backends are unhealthy — unlinking pipe %s to unblock Bazel\n", bb.bepBinPath)
+		if err := syscall.Unlink(bb.bepBinPath); err != nil && !os.IsNotExist(err) {
+			fmt.Fprintf(os.Stderr, "failed to unlink BES pipe %s: %v\n", bb.bepBinPath, err)
+		}
+	})
+}
+
 func (bb *besPipe) streamBesEvents(ctx context.Context, conn *os.File) error {
 	reader := bufio.NewReader(conn)
 
 	go func() {
 		<-ctx.Done()
-		if err := conn.SetReadDeadline(time.Now().Add(gracePeriodDuration)); err != nil {
-			fmt.Fprintf(os.Stderr, "failed to set read deadline after context done: %s\n", err.Error())
-		}
+		conn.SetReadDeadline(time.Now().Add(gracePeriodDuration))
 	}()
 
-	// Manually manage a sequence ID for the events
 	seqId := int64(0)
-
 	besEventGlobalTimeout := time.After(besEventGlobalTimeoutDuration)
+
 	for {
 		event := buildeventstream.BuildEvent{}
-
 		opts := protodelim.UnmarshalOptions{
-			MaxSize: 32 * 1024 * 1024, // 32 MB max; we have observed 17 MB BES events in the wild
+			MaxSize: 32 * 1024 * 1024, // 32 MB max
 		}
 
 		if err := opts.UnmarshalFrom(reader, &event); err != nil {
@@ -220,16 +251,13 @@ func (bb *besPipe) streamBesEvents(ctx context.Context, conn *os.File) error {
 				case <-besEventGlobalTimeout:
 					return fmt.Errorf("timeout reached while waiting for BES events")
 				case <-time.After(besEventThrottleDuration):
-					// throttle the reading of the BES file when no new data is available
 					continue
 				}
 			}
 			return fmt.Errorf("failed to parse BES event: %w", err)
 		}
 
-		// Reset the global timeout on each received event
 		besEventGlobalTimeout = time.After(besEventGlobalTimeoutDuration)
-
 		seqId++
 
 		if err := bb.publishBesEvent(seqId, &event); err != nil {
@@ -241,9 +269,7 @@ func (bb *besPipe) streamBesEvents(ctx context.Context, conn *os.File) error {
 		}
 	}
 
-	// Clear read deadline
 	conn.SetReadDeadline(time.Time{})
-
 	return nil
 }
 
@@ -269,7 +295,6 @@ func (bb *besPipe) publishBesEvent(seqId int64, event *buildeventstream.BuildEve
 			return fmt.Errorf("failed to marshal BES event: %w", err)
 		}
 
-		// Wrap the event in the gRPC message
 		grpcEvent := &buildv1.PublishBuildToolEventStreamRequest{
 			OrderedBuildEvent: &buildv1.OrderedBuildEvent{
 				SequenceNumber: seqId,
@@ -285,14 +310,36 @@ func (bb *besPipe) publishBesEvent(seqId int64, event *buildeventstream.BuildEve
 		}
 
 		for _, p := range bb.besProxies {
-			eg.Go(
-				func() error {
-					if err := p.Send(grpcEvent); err != nil {
+			p := p // capture
+			eg.Go(func() error {
+				if !p.IsHealthy() {
+					return nil
+				}
+
+				// Channel for Send result
+				sendCh := make(chan error, 1)
+
+				// Run Send in goroutine
+				go func() {
+					err := p.Send(grpcEvent)
+					sendCh <- err
+				}()
+
+				// Wait for Send or timeout
+				select {
+				case err := <-sendCh:
+					if err != nil {
 						fmt.Fprintf(os.Stderr, "Error sending BES event to %v: %s\n", p.Host(), err.Error())
+						bb.maybeAbortPipeBecauseNoHealthyBackends()
 					}
 					return nil
-				},
-			)
+				case <-time.After(sendTimeout):
+					fmt.Fprintf(os.Stderr, "Timeout sending BES event to %v: marking unhealthy\n", p.Host())
+					p.MarkUnhealthy()
+					bb.maybeAbortPipeBecauseNoHealthyBackends()
+					return nil
+				}
+			})
 		}
 	}
 
@@ -305,8 +352,6 @@ func (bb *besPipe) Args() []string {
 		bb.bepBinPath,
 	}
 
-	// Also add wait_for_upload_complete flag if the bes pipe was explicitly requested.
-	// NOTE: this is explicitly not the default behavior to avoid breaking changes in bazel6
 	if os.Getenv("ASPECT_BEP_USE_PIPE") != "" {
 		args = append(args, "--build_event_binary_file_upload_mode=wait_for_upload_complete")
 	}
@@ -326,6 +371,6 @@ func (bb *besPipe) Errors() []error {
 
 func (bb *besPipe) GracefulStop() {
 	bb.wg.Wait()
-
+	// Best-effort cleanup
 	os.Remove(bb.bepBinPath)
 }
diff --git a/pkg/plugin/system/besproxy/bes_proxy.go b/pkg/plugin/system/besproxy/bes_proxy.go
@@ -37,7 +37,8 @@ type BESProxy interface {
 	PublishBuildToolEventStream(ctx context.Context, opts ...grpc.CallOption) error
 	PublishLifecycleEvent(ctx context.Context, req *buildv1.PublishLifecycleEventRequest, opts ...grpc.CallOption) (*emptypb.Empty, error)
 	StreamCreated() bool
-	Healthy() bool
+	IsHealthy() bool
+	MarkUnhealthy()
 	Recv() (*buildv1.PublishBuildToolEventStreamResponse, error)
 	Send(req *buildv1.PublishBuildToolEventStreamRequest) error
 }
@@ -142,6 +143,11 @@ func (bp *besProxy) trackError(err error) error {
 	return err
 }
 
-func (bp *besProxy) Healthy() bool {
+func (bp *besProxy) IsHealthy() bool {
 	return bp.hadError < 5 && bp.stream != nil
 }
+
+func (bp *besProxy) MarkUnhealthy() {
+	bp.hadError = 5
+	fmt.Printf("stream to %s is marked unhealthy, taking out of rotation.", bp.host)
+}
