Fix forwardingStreamHandler errgroup race

Commit 5b5db75,
"Add generic gRPC stream forwarding (#306)",
introduced a flakiness in the grpc_test due to calling
`errgroup.Go(func() error { return err })` when an error had occurred in
a go routine that was asynchrounous to the errgroup.Wait() call. This
could lead to `errgroup.Go()` being called just when `errgroup.Wait()`
was about to return, which is not allowed.

Author: moroten

pkg/grpc/BUILD.bazel

@@ -75,7 +75,6 @@ go_library(
         "@org_golang_google_protobuf//proto",
         "@org_golang_google_protobuf//reflect/protoreflect",
         "@org_golang_google_protobuf//types/known/emptypb",
-        "@org_golang_x_sync//errgroup",
         "@org_golang_x_sync//semaphore",
     ] + select({
         "@rules_go//go/platform:android": [

pkg/grpc/forwarding_stream_handler.go

@@ -1,9 +1,10 @@
 package grpc
 
 import (
+	"context"
 	"io"
 
-	"golang.org/x/sync/errgroup"
+	"github.com/buildbarn/bb-storage/pkg/program"
 	"google.golang.org/grpc"
 	"google.golang.org/protobuf/types/known/emptypb"
 )
@@ -38,42 +39,40 @@ func (s *forwardingStreamHandler) HandleStream(srv any, incomingStream grpc.Serv
 		ServerStreams: true,
 		ClientStreams: true,
 	}
-	group, groupCtx := errgroup.WithContext(incomingStream.Context())
-	group.Go(func() error {
-		// groupCtx is guaranteed to be canceled before returning from this method, so outgoingStream will not leak resources.
-		outgoingStream, err := s.backend.NewStream(groupCtx, &desc, method)
+
+	return program.RunLocal(incomingStream.Context(), func(ctx context.Context, siblingsGroup, dependenciesGroup program.Group) error {
+		// ctx is guaranteed to be canceled when returning from this function,
+		// so outgoingStream will not leak resources.
+		outgoingStream, err := s.backend.NewStream(ctx, &desc, method)
 		if err != nil {
 			return err
 		}
-		// Avoid group.Go because incomingStream.RecvMsg might block returning
-		// an error from the outgoingStream and getting the context for
-		// incomingStream canceled.
-		go func() {
+
+		// The only way to cancel a blocking incomingStream.RecvMsg is to return
+		// from this method. Therefore, an error from outgoingStream.RecvMsg
+		// needs to be returned without waiting for incomingStream.RecvMsg.
+		siblingsGroup.GoAsync(func() error {
 			for {
 				msg := &emptypb.Empty{}
 				if err := incomingStream.RecvMsg(msg); err != nil {
 					if err == io.EOF {
 						// Let's continue to receive on outgoingStream, so don't
 						// cancel grouptCtx.
 						outgoingStream.CloseSend()
-						return
+						return nil
 					}
-					// Cancel groupCtx immediately.
-					group.Go(func() error { return err })
-					return
+					return err
 				}
 				if err := outgoingStream.SendMsg(msg); err != nil {
 					if err == io.EOF {
 						// The error will be returned by outgoingStream.RecvMsg(),
 						// no need to cancel groupCtx now.
-						return
+						return nil
 					}
-					// Cancel groupCtx immediately.
-					group.Go(func() error { return err })
-					return
+					return err
 				}
 			}
-		}()
+		})
 
 		for {
 			msg := &emptypb.Empty{}
@@ -88,5 +87,4 @@ func (s *forwardingStreamHandler) HandleStream(srv any, incomingStream grpc.Serv
 			}
 		}
 	})
-	return group.Wait()
 }