test(storage): test Copier.Run's token propagation + refactor int test (#6882)

BrennaEpp · tritone · gcf-merge-on-green[bot] · web-flow · commit c987e969e4f9 · 2022-11-03T13:01:14.000-04:00
Co-authored-by: Chris Cotter &lt;cjcotter@google.com&gt;
Co-authored-by: gcf-merge-on-green[bot] &lt;60162190+gcf-merge-on-green[bot]@users.noreply.github.com&gt;
diff --git a/storage/client.go b/storage/client.go
@@ -317,10 +317,11 @@ type destinationObject struct {
 }
 
 type rewriteObjectRequest struct {
-	srcObject     sourceObject
-	dstObject     destinationObject
-	predefinedACL string
-	token         string
+	srcObject                sourceObject
+	dstObject                destinationObject
+	predefinedACL            string
+	token                    string
+	maxBytesRewrittenPerCall int64
 }
 
 type rewriteObjectResponse struct {
diff --git a/storage/copy.go b/storage/copy.go
@@ -69,6 +69,15 @@ type Copier struct {
 	DestinationKMSKeyName string
 
 	dst, src *ObjectHandle
+
+	// The maximum number of bytes that will be rewritten per rewrite request.
+	// Most callers shouldn't need to specify this parameter - it is primarily
+	// in place to support testing. If specified the value must be an integral
+	// multiple of 1 MiB (1048576). Also, this only applies to requests where
+	// the source and destination span locations and/or storage classes. Finally,
+	// this value must not change across rewrite calls else you'll get an error
+	// that the `rewriteToken` is invalid.
+	maxBytesRewrittenPerCall int64
 }
 
 // Run performs the copy.
@@ -108,8 +117,9 @@ func (c *Copier) Run(ctx context.Context) (attrs *ObjectAttrs, err error) {
 			encryptionKey: c.dst.encryptionKey,
 			keyName:       c.DestinationKMSKeyName,
 		},
-		predefinedACL: c.PredefinedACL,
-		token:         c.RewriteToken,
+		predefinedACL:            c.PredefinedACL,
+		token:                    c.RewriteToken,
+		maxBytesRewrittenPerCall: c.maxBytesRewrittenPerCall,
 	}
 
 	isIdempotent := c.dst.conds != nil && (c.dst.conds.GenerationMatch != 0 || c.dst.conds.DoesNotExist)
diff --git a/storage/grpc_client.go b/storage/grpc_client.go
@@ -822,6 +822,9 @@ func (c *grpcStorageClient) RewriteObject(ctx context.Context, req *rewriteObjec
 		call.CopySourceEncryptionKeyBytes = srcParams.GetEncryptionKeyBytes()
 		call.CopySourceEncryptionKeySha256Bytes = srcParams.GetEncryptionKeySha256Bytes()
 	}
+
+	call.MaxBytesRewrittenPerCall = req.maxBytesRewrittenPerCall
+
 	var res *storagepb.RewriteResponse
 	var err error
 
diff --git a/storage/http_client.go b/storage/http_client.go
@@ -747,6 +747,11 @@ func (c *httpStorageClient) RewriteObject(ctx context.Context, req *rewriteObjec
 	if err := setEncryptionHeaders(call.Header(), req.srcObject.encryptionKey, true); err != nil {
 		return nil, err
 	}
+
+	if req.maxBytesRewrittenPerCall != 0 {
+		call.MaxBytesRewrittenPerCall(req.maxBytesRewrittenPerCall)
+	}
+
 	var res *raw.RewriteResponse
 	var err error
 	setClientHeader(call.Header())
diff --git a/storage/integration_test.go b/storage/integration_test.go
@@ -1454,33 +1454,6 @@ func TestIntegration_Objects(t *testing.T) {
 			t.Errorf("Object %v is older than its containing bucket, %v", o, bAttrs)
 		}
 
-		// Test object copy.
-		copyName := "copy-" + objName
-		copyObj, err := bkt.Object(copyName).CopierFrom(bkt.Object(objName)).Run(ctx)
-		if err != nil {
-			t.Errorf("Copier.Run failed with %v", err)
-		} else if !namesEqual(copyObj, newBucketName, copyName) {
-			t.Errorf("Copy object bucket, name: got %q.%q, want %q.%q",
-				copyObj.Bucket, copyObj.Name, newBucketName, copyName)
-		}
-
-		// Copying with attributes.
-		const contentEncoding = "identity"
-		copier := bkt.Object(copyName).CopierFrom(bkt.Object(objName))
-		copier.ContentEncoding = contentEncoding
-		copyObj, err = copier.Run(ctx)
-		if err != nil {
-			t.Errorf("Copier.Run failed with %v", err)
-		} else {
-			if !namesEqual(copyObj, newBucketName, copyName) {
-				t.Errorf("Copy object bucket, name: got %q.%q, want %q.%q",
-					copyObj.Bucket, copyObj.Name, newBucketName, copyName)
-			}
-			if copyObj.ContentEncoding != contentEncoding {
-				t.Errorf("Copy ContentEncoding: got %q, want %q", copyObj.ContentEncoding, contentEncoding)
-			}
-		}
-
 		objectHandle := bkt.Object(objName)
 
 		// Test UpdateAttrs.
@@ -1598,17 +1571,6 @@ func TestIntegration_Objects(t *testing.T) {
 			t.Error("Close expected an error, found none")
 		}
 
-		// Test deleting the copy object.
-		h.mustDeleteObject(bkt.Object(copyName))
-		// Deleting it a second time should return ErrObjectNotExist.
-		if err := bkt.Object(copyName).Delete(ctx); err != ErrObjectNotExist {
-			t.Errorf("second deletion of %v = %v; want ErrObjectNotExist", copyName, err)
-		}
-		_, err = bkt.Object(copyName).Attrs(ctx)
-		if err != ErrObjectNotExist {
-			t.Errorf("Copy is expected to be deleted, stat errored with %v", err)
-		}
-
 		// Test object composition.
 		var compSrcs []*ObjectHandle
 		var wantContents []byte
@@ -1650,6 +1612,150 @@ func TestIntegration_Objects(t *testing.T) {
 	})
 }
 
+func TestIntegration_Copy(t *testing.T) {
+	multiTransportTest(context.Background(), t, func(t *testing.T, ctx context.Context, bucket string, prefix string, client *Client) {
+		bucketFrom := client.Bucket(bucket)
+		bucketInSameRegion := client.Bucket(prefix + uidSpace.New())
+		bucketInDifferentRegion := client.Bucket(prefix + uidSpace.New())
+
+		// Create new bucket
+		if err := bucketInSameRegion.Create(ctx, testutil.ProjID(), nil); err != nil {
+			t.Fatalf("bucket.Create: %v", err)
+		}
+		defer bucketInSameRegion.Delete(ctx)
+
+		// Create new bucket
+		if err := bucketInDifferentRegion.Create(ctx, testutil.ProjID(), &BucketAttrs{Location: "NORTHAMERICA-NORTHEAST2"}); err != nil {
+			t.Fatalf("bucket.Create: %v", err)
+		}
+		defer bucketInDifferentRegion.Delete(ctx)
+
+		// We use a larger object size to be able to trigger multiple rewrite calls
+		minObjectSize := 2500000 // 2.5 Mb
+		obj := bucketFrom.Object("copy-object-original" + uidSpace.New())
+
+		// Create an object to copy from
+		w := obj.NewWriter(ctx)
+		c := randomContents()
+		for written := 0; written < minObjectSize; {
+			n, err := w.Write(c)
+			if err != nil {
+				t.Fatalf("w.Write: %v", err)
+			}
+			written += n
+		}
+		if err := w.Close(); err != nil {
+			t.Fatalf("w.Close: %v", err)
+		}
+
+		defer func() {
+			if err := obj.Delete(ctx); err != nil {
+				t.Errorf("obj.Delete: %v", err)
+			}
+		}()
+
+		attrs, err := obj.Attrs(ctx)
+		if err != nil {
+			t.Fatalf("obj.Attrs: %v", err)
+		}
+
+		crc32c := attrs.CRC32C
+
+		type copierAttrs struct {
+			contentEncoding string
+			maxBytesPerCall int64
+		}
+
+		for _, test := range []struct {
+			desc                    string
+			toObj                   string
+			toBucket                *BucketHandle
+			copierAttrs             *copierAttrs
+			numExpectedRewriteCalls int
+		}{
+			{
+				desc:                    "copy within bucket",
+				toObj:                   "copy-within-bucket",
+				toBucket:                bucketFrom,
+				numExpectedRewriteCalls: 1,
+			},
+			{
+				desc:                    "copy to new bucket",
+				toObj:                   "copy-new-bucket",
+				toBucket:                bucketInSameRegion,
+				numExpectedRewriteCalls: 1,
+			},
+			{
+				desc:                    "copy with attributes",
+				toObj:                   "copy-with-attributes",
+				toBucket:                bucketInSameRegion,
+				copierAttrs:             &copierAttrs{contentEncoding: "identity"},
+				numExpectedRewriteCalls: 1,
+			},
+			{
+				// this test should trigger multiple re-write calls and may fail
+				// with a rate limit error if those calls are stuck in an infinite loop
+				desc:                    "copy to new region",
+				toObj:                   "copy-new-region",
+				toBucket:                bucketInDifferentRegion,
+				copierAttrs:             &copierAttrs{maxBytesPerCall: 1048576},
+				numExpectedRewriteCalls: 3,
+			},
+		} {
+			t.Run(test.desc, func(t *testing.T) {
+				copyObj := test.toBucket.Object(test.toObj)
+				copier := copyObj.CopierFrom(obj)
+
+				if attrs := test.copierAttrs; attrs != nil {
+					if attrs.contentEncoding != "" {
+						copier.ContentEncoding = attrs.contentEncoding
+					}
+					if attrs.maxBytesPerCall != 0 {
+						copier.maxBytesRewrittenPerCall = attrs.maxBytesPerCall
+					}
+				}
+
+				rewriteCallsCount := 0
+				copier.ProgressFunc = func(_, _ uint64) {
+					rewriteCallsCount++
+				}
+
+				attrs, err = copier.Run(ctx)
+				if err != nil {
+					t.Fatalf("Copier.Run failed with %v", err)
+				}
+				defer func() {
+					if err := copyObj.Delete(ctx); err != nil {
+						t.Errorf("copyObj.Delete: %v", err)
+					}
+				}()
+
+				// Check copied object is in the correct bucket with the correct name
+				if attrs.Bucket != test.toBucket.name || attrs.Name != test.toObj {
+					t.Errorf("unexpected copy behaviour: got: %s in bucket %s, want: %s in bucket %s", attrs.Name, attrs.Bucket, attrs.Name, test.toBucket.name)
+				}
+
+				// Check attrs
+				if test.copierAttrs != nil {
+					if attrs.ContentEncoding != test.copierAttrs.contentEncoding {
+						t.Errorf("unexpected ContentEncoding; got: %s, want: %s", attrs.ContentEncoding, test.copierAttrs.contentEncoding)
+					}
+				}
+
+				// Check the copied contents
+				if attrs.CRC32C != crc32c {
+					t.Errorf("mismatching checksum: got %v, want %v", attrs.CRC32C, crc32c)
+				}
+
+				// Check that the number of requests made is as expected
+				if rewriteCallsCount != test.numExpectedRewriteCalls {
+					t.Errorf("unexpected number of rewrite calls: got %v, want %v", rewriteCallsCount, test.numExpectedRewriteCalls)
+				}
+			})
+		}
+	})
+}
+
 func TestIntegration_Encoding(t *testing.T) {
 	ctx := context.Background()
 	client := testConfig(ctx, t)
@@ -4777,10 +4883,6 @@ func putURL(url string, headers map[string][]string, payload io.Reader) ([]byte,
 	return bytes, nil
 }
 
-func namesEqual(obj *ObjectAttrs, bucketName, objectName string) bool {
-	return obj.Bucket == bucketName && obj.Name == objectName
-}
-
 func keyFileEmail(filename string) (string, error) {
 	bytes, err := ioutil.ReadFile(filename)
 	if err != nil {

Original file line number	Diff line number	Diff line change
`@@ -317,10 +317,11 @@ type destinationObject struct {`
`317`	`317`	`}`
`318`	`318`
`319`	`319`	`type rewriteObjectRequest struct {`
`320`		`- srcObject sourceObject`
`321`		`- dstObject destinationObject`
`322`		`- predefinedACL string`
`323`		`- token string`
	`320`	`+ srcObject sourceObject`
	`321`	`+ dstObject destinationObject`
	`322`	`+ predefinedACL string`
	`323`	`+ token string`
	`324`	`+ maxBytesRewrittenPerCall int64`
`324`	`325`	`}`
`325`	`326`
`326`	`327`	`type rewriteObjectResponse struct {`
Original file line number	Diff line number	Diff line change
`@@ -822,6 +822,9 @@ func (c grpcStorageClient) RewriteObject(ctx context.Context, req rewriteObjec`
`822`	`822`	`call.CopySourceEncryptionKeyBytes = srcParams.GetEncryptionKeyBytes()`
`823`	`823`	`call.CopySourceEncryptionKeySha256Bytes = srcParams.GetEncryptionKeySha256Bytes()`
`824`	`824`	`}`
	`825`	`+`
	`826`	`+ call.MaxBytesRewrittenPerCall = req.maxBytesRewrittenPerCall`
	`827`	`+`
`825`	`828`	`var res *storagepb.RewriteResponse`
`826`	`829`	`var err error`
`827`	`830`