fix(storage): ZB Reader redirect support (#12703)

* fix(storage): ZB Reader redirect support

* do single retry for redirect

* add mid-stream redirect for aborted

* clean up test comments

Author: tritone

storage/grpc_client.go

@@ -1628,7 +1628,6 @@ func (c *grpcStorageClient) NewRangeReader(ctx context.Context, params *newRange
 	req := &storagepb.BidiReadObjectRequest{
 		ReadObjectSpec: spec,
 	}
-	ctx = gax.InsertMetadataIntoOutgoingContext(ctx, contextMetadataFromBidiReadObject(req)...)
 
 	// Define a function that initiates a Read with offset and length, assuming
 	// we have already read seen bytes.
@@ -1660,28 +1659,53 @@ func (c *grpcStorageClient) NewRangeReader(ctx context.Context, params *newRange
 		var decoder *readResponseDecoder
 
 		err = run(cc, func(ctx context.Context) error {
-			stream, err = c.raw.BidiReadObject(ctx, s.gax...)
-			if err != nil {
-				return err
-			}
-			if err := stream.Send(req); err != nil {
-				return err
+			var databufs mem.BufferSlice
+			openAndSendReq := func() error {
+				databufs = mem.BufferSlice{}
+
+				// Insert context metadata, including routing token if this is a retry
+				// for a redirect.
+				mdCtx := gax.InsertMetadataIntoOutgoingContext(ctx, contextMetadataFromBidiReadObject(req)...)
+				stream, err = c.raw.BidiReadObject(mdCtx, s.gax...)
+				if err != nil {
+					return err
+				}
+				if err := stream.Send(req); err != nil {
+					return err
+				}
+				// Oneshot reads can close the client->server side immediately.
+				if err := stream.CloseSend(); err != nil {
+					return err
+				}
+
+				// Receive the message into databuf as a wire-encoded message so we can
+				// use a custom decoder to avoid an extra copy at the protobuf layer.
+				return stream.RecvMsg(&databufs)
 			}
-			// Oneshot reads can close the client->server side immediately.
-			if err := stream.CloseSend(); err != nil {
-				return err
+
+			err := openAndSendReq()
+
+			// We might get a redirect error here for an out-of-region request.
+			// Add the routing token and read handle to the request and do one
+			// retry.
+			if st, ok := status.FromError(err); ok && st.Code() == codes.Aborted {
+				for _, d := range st.Details() {
+					if e, ok := d.(*storagepb.BidiReadObjectRedirectedError); ok {
+						req.ReadObjectSpec.ReadHandle = e.GetReadHandle()
+						req.ReadObjectSpec.RoutingToken = e.RoutingToken
+						err = openAndSendReq()
+						break
+					}
+				}
 			}
 
-			// Receive the message into databuf as a wire-encoded message so we can
-			// use a custom decoder to avoid an extra copy at the protobuf layer.
-			databufs := mem.BufferSlice{}
-			err := stream.RecvMsg(&databufs)
 			// These types of errors show up on the RecvMsg call, rather than the
 			// initialization of the stream via BidiReadObject above.
 			if s, ok := status.FromError(err); ok && s.Code() == codes.NotFound {
-				return formatObjectErr(err)
+				err = formatObjectErr(err)
 			}
 			if err != nil {
+				databufs.Free()
 				return err
 			}
 			// Use a custom decoder that uses protobuf unmarshalling for all
@@ -2093,7 +2117,9 @@ func (r *gRPCReader) Close() error {
 func (r *gRPCReader) recv() error {
 	databufs := mem.BufferSlice{}
 	err := r.stream.RecvMsg(&databufs)
-	if err != nil && r.settings.retry.runShouldRetry(err) {
+	// If we get a mid-stream error on a recv call, reopen the stream.
+	// ABORTED could indicate a redirect so should also trigger a reopen.
+	if err != nil && (r.settings.retry.runShouldRetry(err) || status.Code(err) == codes.Aborted) {
 		// This will "close" the existing stream and immediately attempt to
 		// reopen the stream, but will backoff if further attempts are necessary.
 		// Reopening the stream Recvs the first message, so if retrying is

storage/integration_test.go

@@ -46,7 +46,6 @@ import (
 
 	"cloud.google.com/go/auth"
 	"cloud.google.com/go/auth/credentials"
-	"cloud.google.com/go/compute/metadata"
 	"cloud.google.com/go/httpreplay"
 	"cloud.google.com/go/iam"
 	"cloud.google.com/go/iam/apiv1/iampb"
@@ -6933,6 +6932,43 @@ func TestIntegration_OTelTracing(t *testing.T) {
 	})
 }
 
+// Simple integration test for a zonal bucket read via storage.Reader.
+// Will test out-of-region redirect flow if it's run from outside of us-west4.
+func TestIntegration_ZonalRead(t *testing.T) {
+	multiTransportTest(skipAllButBidi(context.Background(), "zonal bucket test"), t, func(t *testing.T, ctx context.Context, _ string, prefix string, client *Client) {
+		h := testHelper{t}
+		bucketName := prefix + uidSpace.New()
+		bkt := client.Bucket(bucketName)
+
+		h.mustCreateZonalBucket(bkt, testutil.ProjID())
+		defer h.mustDeleteBucket(bkt)
+
+		objName := "bidi-test-obj"
+		obj := bkt.Object(objName)
+		w := obj.NewWriter(ctx)
+		w.Append = true
+		h.mustWrite(w, randomBytes3MiB)
+		defer obj.Delete(ctx)
+
+		r, err := obj.NewReader(ctx)
+		if err != nil {
+			t.Fatalf("NewReader: %v", err)
+		}
+
+		b, err := io.ReadAll(r)
+		if err != nil {
+			t.Fatalf("reading: %v", err)
+		}
+		if !bytes.Equal(randomBytes3MiB, b) {
+			t.Errorf("download data does not match; got %v bytes, want %v bytes", len(b), len(randomBytes3MiB))
+		}
+
+		if err := r.Close(); err != nil {
+			t.Errorf("closing: %v", err)
+		}
+	}, experimental.WithGRPCBidiReads())
+}
+
 // openTelemetryTestExporter is a test utility exporter. It should be created
 // with NewopenTelemetryTestExporter.
 type openTelemetryTestExporter struct {
@@ -7207,13 +7243,13 @@ func (h testHelper) mustCreate(b *BucketHandle, projID string, attrs *BucketAttr
 func (h testHelper) mustCreateZonalBucket(b *BucketHandle, projID string) {
 	h.t.Helper()
 
-	// Create a bucket in the same zone as the test VM.
-	zone, err := metadata.ZoneWithContext(context.Background())
-	if err != nil {
-		h.t.Fatalf("could not determine VM zone: %v", err)
-	}
-	region := strings.Join(strings.Split(zone, "-")[:2], "-")
-	h.mustCreate(b, testutil.ProjID(), &BucketAttrs{
+	// Create a zonal bucket in us-west4-a.
+	// Another zone/region can be subbed in if you want to run elsewhere because
+	// of quota reasons or to run in the same region as your VM.
+	zone := "us-west4-a"
+	region := "us-west4"
+
+	h.mustCreate(b, projID, &BucketAttrs{
 		Location: region,
 		CustomPlacementConfig: &CustomPlacementConfig{
 			DataLocations: []string{zone},