|
82 | 82 | import com.google.cloud.storage.StorageException; |
83 | 83 | import com.google.cloud.storage.StorageOptions; |
84 | 84 | import com.google.cloud.storage.StorageRoles; |
| 85 | +import com.google.cloud.storage.spi.StorageRpcFactory; |
| 86 | +import com.google.cloud.storage.spi.v1.StorageRpc; |
| 87 | +import com.google.cloud.storage.spi.v1.StorageRpc.Option; |
85 | 88 | import com.google.cloud.storage.testing.RemoteStorageHelper; |
86 | 89 | import com.google.common.collect.ImmutableList; |
87 | 90 | import com.google.common.collect.ImmutableMap; |
|
90 | 93 | import com.google.common.collect.Lists; |
91 | 94 | import com.google.common.io.BaseEncoding; |
92 | 95 | import com.google.common.io.ByteStreams; |
| 96 | +import com.google.common.reflect.AbstractInvocationHandler; |
| 97 | +import com.google.common.reflect.Reflection; |
93 | 98 | import com.google.iam.v1.Binding; |
94 | 99 | import com.google.iam.v1.IAMPolicyGrpc; |
95 | 100 | import com.google.iam.v1.SetIamPolicyRequest; |
|
107 | 112 | import java.io.FileOutputStream; |
108 | 113 | import java.io.IOException; |
109 | 114 | import java.io.InputStream; |
| 115 | +import java.lang.reflect.Method; |
110 | 116 | import java.net.URL; |
111 | 117 | import java.net.URLConnection; |
112 | 118 | import java.nio.ByteBuffer; |
| 119 | +import java.nio.charset.StandardCharsets; |
113 | 120 | import java.nio.file.Files; |
114 | 121 | import java.nio.file.Path; |
115 | 122 | import java.security.Key; |
|
125 | 132 | import java.util.Set; |
126 | 133 | import java.util.concurrent.ExecutionException; |
127 | 134 | import java.util.concurrent.TimeUnit; |
| 135 | +import java.util.concurrent.atomic.AtomicBoolean; |
128 | 136 | import java.util.logging.Level; |
129 | 137 | import java.util.logging.Logger; |
130 | 138 | import java.util.zip.GZIPInputStream; |
|
139 | 147 | import org.apache.http.impl.conn.PoolingHttpClientConnectionManager; |
140 | 148 | import org.junit.AfterClass; |
141 | 149 | import org.junit.BeforeClass; |
| 150 | +import org.junit.Rule; |
142 | 151 | import org.junit.Test; |
| 152 | +import org.junit.rules.TestName; |
| 153 | +import org.threeten.bp.Clock; |
| 154 | +import org.threeten.bp.Instant; |
| 155 | +import org.threeten.bp.ZoneId; |
| 156 | +import org.threeten.bp.ZoneOffset; |
| 157 | +import org.threeten.bp.format.DateTimeFormatter; |
143 | 158 |
|
144 | 159 | public class ITStorageTest { |
145 | 160 |
|
@@ -200,6 +215,8 @@ public class ITStorageTest { |
200 | 215 | private static final ImmutableList<LifecycleRule> LIFECYCLE_RULES = |
201 | 216 | ImmutableList.of(LIFECYCLE_RULE_1, LIFECYCLE_RULE_2); |
202 | 217 |
|
| 218 | +@Rule public final TestName testName = new TestName(); |
| 219 | + |
203 | 220 | @BeforeClass |
204 | 221 | public static void beforeClass() throws IOException { |
205 | 222 | remoteStorageHelper = RemoteStorageHelper.create(); |
@@ -3813,4 +3830,125 @@ public void testWriterWithKmsKeyName() throws IOException { |
3813 | 3830 | assertThat(blob.getKmsKeyName()).isNotNull(); |
3814 | 3831 | assertThat(storage.delete(BUCKET, blobName)).isTrue(); |
3815 | 3832 | } |
| 3833 | + |
| 3834 | +@Test |
| 3835 | +public void blobWriteChannel_handlesRecoveryOnLastChunkWhenGenerationIsPresent_multipleChunks() |
| 3836 | +throws IOException { |
| 3837 | +int _2MiB = 256 * 1024; |
| 3838 | +int contentSize = 292_617; |
| 3839 | + |
| 3840 | +blobWriteChannel_handlesRecoveryOnLastChunkWhenGenerationIsPresent(_2MiB, contentSize); |
| 3841 | +} |
| 3842 | + |
| 3843 | +@Test |
| 3844 | +public void blobWriteChannel_handlesRecoveryOnLastChunkWhenGenerationIsPresent_singleChunk() |
| 3845 | +throws IOException { |
| 3846 | +int _4MiB = 256 * 1024 * 2; |
| 3847 | +int contentSize = 292_617; |
| 3848 | + |
| 3849 | +blobWriteChannel_handlesRecoveryOnLastChunkWhenGenerationIsPresent(_4MiB, contentSize); |
| 3850 | +} |
| 3851 | + |
| 3852 | +private void blobWriteChannel_handlesRecoveryOnLastChunkWhenGenerationIsPresent( |
| 3853 | +int chunkSize, int contentSize) throws IOException { |
| 3854 | +Instant now = Clock.systemUTC().instant(); |
| 3855 | +DateTimeFormatter formatter = |
| 3856 | +DateTimeFormatter.ISO_LOCAL_DATE_TIME.withZone(ZoneId.from(ZoneOffset.UTC)); |
| 3857 | +String nowString = formatter.format(now); |
| 3858 | + |
| 3859 | +String blobPath = String.format("%s/%s/blob", testName.getMethodName(), nowString); |
| 3860 | +BlobId blobId = BlobId.of(BUCKET, blobPath); |
| 3861 | +BlobInfo blobInfo = BlobInfo.newBuilder(blobId).build(); |
| 3862 | + |
| 3863 | +Random rand = new Random(1234567890); |
| 3864 | +String randString = randString(rand, contentSize); |
| 3865 | +final byte[] randStringBytes = randString.getBytes(StandardCharsets.UTF_8); |
| 3866 | +Storage storage = StorageOptions.getDefaultInstance().getService(); |
| 3867 | +WriteChannel ww = storage.writer(blobInfo); |
| 3868 | +ww.setChunkSize(chunkSize); |
| 3869 | +ww.write(ByteBuffer.wrap(randStringBytes)); |
| 3870 | +ww.close(); |
| 3871 | + |
| 3872 | +Blob blobGen1 = storage.get(blobId); |
| 3873 | + |
| 3874 | +final AtomicBoolean exceptionThrown = new AtomicBoolean(false); |
| 3875 | + |
| 3876 | +Storage testStorage = |
| 3877 | +StorageOptions.newBuilder() |
| 3878 | +.setServiceRpcFactory( |
| 3879 | +new StorageRpcFactory() { |
| 3880 | +/** |
| 3881 | +* Here we're creating a proxy of StorageRpc where we can delegate all calls to |
| 3882 | +* the normal implementation, except in the case of {@link |
| 3883 | +* StorageRpc#writeWithResponse(String, byte[], int, long, int, boolean)} where |
| 3884 | +* {@code lastChunk == true}. We allow the call to execute, but instead of |
| 3885 | +* returning the result we throw an IOException to simulate a prematurely close |
| 3886 | +* connection. This behavior is to ensure appropriate handling of a completed |
| 3887 | +* upload where the ACK wasn't received. In particular, if an upload is initiated |
| 3888 | +* against an object where an {@link Option#IF_GENERATION_MATCH} simply calling |
| 3889 | +* get on an object can result in a 404 because the object that is created while |
| 3890 | +* the BlobWriteChannel is executing will be a new generation. |
| 3891 | +*/ |
| 3892 | +@SuppressWarnings("UnstableApiUsage") |
| 3893 | +@Override |
| 3894 | +public StorageRpc create(final StorageOptions options) { |
| 3895 | +return Reflection.newProxy( |
| 3896 | +StorageRpc.class, |
| 3897 | +new AbstractInvocationHandler() { |
| 3898 | +final StorageRpc delegate = |
| 3899 | +(StorageRpc) StorageOptions.getDefaultInstance().getRpc(); |
| 3900 | + |
| 3901 | +@Override |
| 3902 | +protected Object handleInvocation( |
| 3903 | +Object proxy, Method method, Object[] args) throws Throwable { |
| 3904 | +if ("writeWithResponse".equals(method.getName())) { |
| 3905 | +Object result = method.invoke(delegate, args); |
| 3906 | +boolean lastChunk = (boolean) args[5]; |
| 3907 | +// if we're on the lastChunk simulate a connection failure which |
| 3908 | +// happens after the request was processed but before response could |
| 3909 | +// be received by the client. |
| 3910 | +if (lastChunk) { |
| 3911 | +exceptionThrown.set(true); |
| 3912 | +throw StorageException.translate( |
| 3913 | +new IOException("simulated Connection closed prematurely")); |
| 3914 | +} else { |
| 3915 | +return result; |
| 3916 | +} |
| 3917 | +} |
| 3918 | +return method.invoke(delegate, args); |
| 3919 | +} |
| 3920 | +}); |
| 3921 | +} |
| 3922 | +}) |
| 3923 | +.build() |
| 3924 | +.getService(); |
| 3925 | + |
| 3926 | +try (WriteChannel w = testStorage.writer(blobGen1, BlobWriteOption.generationMatch())) { |
| 3927 | +w.setChunkSize(chunkSize); |
| 3928 | + |
| 3929 | +ByteBuffer buffer = ByteBuffer.wrap(randStringBytes); |
| 3930 | +w.write(buffer); |
| 3931 | +} |
| 3932 | + |
| 3933 | +assertTrue("Expected an exception to be thrown for the last chunk", exceptionThrown.get()); |
| 3934 | + |
| 3935 | +Blob blobGen2 = storage.get(blobId); |
| 3936 | +assertEquals(contentSize, (long) blobGen2.getSize()); |
| 3937 | +assertNotEquals(blobInfo.getGeneration(), blobGen2.getGeneration()); |
| 3938 | +ByteArrayOutputStream actualData = new ByteArrayOutputStream(); |
| 3939 | +blobGen2.downloadTo(actualData); |
| 3940 | +assertArrayEquals(randStringBytes, actualData.toByteArray()); |
| 3941 | +} |
| 3942 | + |
| 3943 | +private static String randString(Random rand, int length) { |
| 3944 | +final StringBuilder sb = new StringBuilder(); |
| 3945 | +while (sb.length() < length) { |
| 3946 | +int i = rand.nextInt('z'); |
| 3947 | +char c = (char) i; |
| 3948 | +if (Character.isLetter(c) || Character.isDigit(c)) { |
| 3949 | +sb.append(c); |
| 3950 | +} |
| 3951 | +} |
| 3952 | +return sb.toString(); |
| 3953 | +} |
3816 | 3954 | } |
0 commit comments