16
16
17
17
package com .google .cloud .storage ;
18
18
19
+ import static com .google .cloud .storage .GrpcUtils .contextWithBucketName ;
20
+
19
21
import com .google .api .core .SettableApiFuture ;
22
+ import com .google .api .gax .grpc .GrpcCallContext ;
23
+ import com .google .api .gax .retrying .ResultRetryAlgorithm ;
24
+ import com .google .api .gax .rpc .ApiStreamObserver ;
25
+ import com .google .api .gax .rpc .BidiStreamingCallable ;
20
26
import com .google .cloud .storage .ChunkSegmenter .ChunkSegment ;
27
+ import com .google .cloud .storage .Conversions .Decoder ;
21
28
import com .google .cloud .storage .Crc32cValue .Crc32cLengthKnown ;
29
+ import com .google .cloud .storage .Retrying .RetryingDependencies ;
22
30
import com .google .cloud .storage .UnbufferedWritableByteChannelSession .UnbufferedWritableByteChannel ;
23
31
import com .google .common .annotations .VisibleForTesting ;
24
32
import com .google .protobuf .ByteString ;
31
39
import java .nio .channels .ClosedChannelException ;
32
40
import java .util .ArrayList ;
33
41
import java .util .List ;
42
+ import java .util .concurrent .Semaphore ;
43
+ import java .util .function .Supplier ;
34
44
import org .checkerframework .checker .nullness .qual .NonNull ;
35
45
36
- final class GapicBidiUnbufferedWritableByteChannel <
37
- RequestFactoryT extends BidiWriteCtx .BidiWriteObjectRequestBuilderFactory >
38
- implements UnbufferedWritableByteChannel {
39
-
46
+ final class GapicBidiUnbufferedWritableByteChannel implements UnbufferedWritableByteChannel {
47
+ private final BidiStreamingCallable <BidiWriteObjectRequest , BidiWriteObjectResponse > write ;
48
+ private final RetryingDependencies deps ;
49
+ private final ResultRetryAlgorithm <?> alg ;
50
+ private final String bucketName ;
51
+ private final Supplier <GrpcCallContext > baseContextSupplier ;
40
52
private final SettableApiFuture <BidiWriteObjectResponse > resultFuture ;
41
53
private final ChunkSegmenter chunkSegmenter ;
42
54
43
- private final BidiWriteCtx <RequestFactoryT > writeCtx ;
44
- private final WriteFlushStrategy . BidiFlusher flusher ;
55
+ private final BidiWriteCtx <BidiResumableWrite > writeCtx ;
56
+ private final BidiObserver responseObserver ;
45
57
58
+ private volatile ApiStreamObserver <BidiWriteObjectRequest > stream ;
46
59
private boolean open = true ;
60
+ private boolean first = true ;
47
61
private boolean finished = false ;
48
62
49
63
GapicBidiUnbufferedWritableByteChannel (
64
+ BidiStreamingCallable <BidiWriteObjectRequest , BidiWriteObjectResponse > write ,
65
+ RetryingDependencies deps ,
66
+ ResultRetryAlgorithm <?> alg ,
50
67
SettableApiFuture <BidiWriteObjectResponse > resultFuture ,
51
68
ChunkSegmenter chunkSegmenter ,
52
- RequestFactoryT requestFactory ,
53
- WriteFlushStrategy .BidiFlusherFactory flusherFactory ) {
69
+ BidiResumableWrite requestFactory ,
70
+ Supplier <GrpcCallContext > baseContextSupplier ) {
71
+ this .write = write ;
72
+ this .deps = deps ;
73
+ this .alg = alg ;
74
+ this .baseContextSupplier = baseContextSupplier ;
75
+ this .bucketName = requestFactory .bucketName ();
54
76
this .resultFuture = resultFuture ;
55
77
this .chunkSegmenter = chunkSegmenter ;
56
78
57
79
this .writeCtx = new BidiWriteCtx <>(requestFactory );
58
- this .flusher =
59
- flusherFactory .newFlusher (
60
- requestFactory .bucketName (), writeCtx .getConfirmedBytes ()::set , resultFuture ::set );
80
+ this .responseObserver = new BidiObserver ();
61
81
}
62
82
63
83
@ Override
64
84
public long write (ByteBuffer [] srcs , int srcsOffset , int srcsLength ) throws IOException {
65
85
return internalWrite (srcs , srcsOffset , srcsLength , false );
66
86
}
67
87
88
+ @ Override
89
+ public long writeAndClose (ByteBuffer [] srcs , int offset , int length ) throws IOException {
90
+ long written = internalWrite (srcs , offset , length , true );
91
+ close ();
92
+ return written ;
93
+ }
94
+
68
95
@ Override
69
96
public boolean isOpen () {
70
97
return open ;
71
98
}
72
99
73
100
@ Override
74
101
public void close () throws IOException {
102
+ if (!open ) {
103
+ return ;
104
+ }
105
+ ApiStreamObserver <BidiWriteObjectRequest > openedStream = openedStream ();
75
106
if (!finished ) {
76
107
BidiWriteObjectRequest message = finishMessage ();
77
108
try {
78
- flusher . close (message );
109
+ openedStream . onNext (message );
79
110
finished = true ;
111
+ openedStream .onCompleted ();
80
112
} catch (RuntimeException e ) {
81
113
resultFuture .setException (e );
82
114
throw e ;
83
115
}
84
116
} else {
85
- flusher . close ( null );
117
+ openedStream . onCompleted ( );
86
118
}
119
+ responseObserver .await ();
87
120
open = false ;
88
121
}
89
122
90
123
@ VisibleForTesting
91
- BidiWriteCtx <RequestFactoryT > getWriteCtx () {
124
+ BidiWriteCtx <BidiResumableWrite > getWriteCtx () {
92
125
return writeCtx ;
93
126
}
94
127
@@ -130,7 +163,8 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, bo
130
163
finished = true ;
131
164
}
132
165
133
- BidiWriteObjectRequest build = builder .build ();
166
+ BidiWriteObjectRequest build = possiblyPairDownBidiRequest (builder , first ).build ();
167
+ first = false ;
134
168
messages .add (build );
135
169
bytesConsumed += contentSize ;
136
170
}
@@ -140,7 +174,7 @@ private long internalWrite(ByteBuffer[] srcs, int srcsOffset, int srcsLength, bo
140
174
}
141
175
142
176
try {
143
- flusher . flush (messages );
177
+ flush (messages );
144
178
} catch (RuntimeException e ) {
145
179
resultFuture .setException (e );
146
180
throw e ;
@@ -162,4 +196,123 @@ private BidiWriteObjectRequest finishMessage() {
162
196
BidiWriteObjectRequest message = b .build ();
163
197
return message ;
164
198
}
199
+
200
+ private ApiStreamObserver <BidiWriteObjectRequest > openedStream () {
201
+ if (stream == null ) {
202
+ synchronized (this ) {
203
+ if (stream == null ) {
204
+ GrpcCallContext internalContext =
205
+ contextWithBucketName (bucketName , baseContextSupplier .get ());
206
+ stream =
207
+ this .write
208
+ .withDefaultCallContext (internalContext )
209
+ .bidiStreamingCall (responseObserver );
210
+ responseObserver .sem .drainPermits ();
211
+ }
212
+ }
213
+ }
214
+ return stream ;
215
+ }
216
+
217
+ private void flush (@ NonNull List <BidiWriteObjectRequest > segments ) {
218
+ Retrying .run (
219
+ deps ,
220
+ alg ,
221
+ () -> {
222
+ try {
223
+ ApiStreamObserver <BidiWriteObjectRequest > opened = openedStream ();
224
+ for (BidiWriteObjectRequest message : segments ) {
225
+ opened .onNext (message );
226
+ }
227
+ if (!finished ) {
228
+ BidiWriteObjectRequest message =
229
+ BidiWriteObjectRequest .newBuilder ().setFlush (true ).setStateLookup (true ).build ();
230
+ opened .onNext (message );
231
+ }
232
+ responseObserver .await ();
233
+ return null ;
234
+ } catch (Exception e ) {
235
+ stream = null ;
236
+ first = true ;
237
+ throw e ;
238
+ }
239
+ },
240
+ Decoder .identity ());
241
+ }
242
+
243
+ private static BidiWriteObjectRequest .Builder possiblyPairDownBidiRequest (
244
+ BidiWriteObjectRequest .Builder b , boolean firstMessageOfStream ) {
245
+ if (firstMessageOfStream && b .getWriteOffset () == 0 ) {
246
+ return b ;
247
+ }
248
+
249
+ if (!firstMessageOfStream ) {
250
+ b .clearUploadId ();
251
+ }
252
+
253
+ if (b .getWriteOffset () > 0 ) {
254
+ b .clearWriteObjectSpec ();
255
+ }
256
+
257
+ if (b .getWriteOffset () > 0 && !b .getFinishWrite ()) {
258
+ b .clearObjectChecksums ();
259
+ }
260
+ return b ;
261
+ }
262
+
263
+ private class BidiObserver implements ApiStreamObserver <BidiWriteObjectResponse > {
264
+
265
+ private final Semaphore sem ;
266
+ private volatile BidiWriteObjectResponse last ;
267
+ private volatile RuntimeException previousError ;
268
+
269
+ private BidiObserver () {
270
+ this .sem = new Semaphore (0 );
271
+ }
272
+
273
+ @ Override
274
+ public void onNext (BidiWriteObjectResponse value ) {
275
+ // incremental update
276
+ if (value .hasPersistedSize ()) {
277
+ writeCtx .getConfirmedBytes ().set ((value .getPersistedSize ()));
278
+ } else if (value .hasResource ()) {
279
+ writeCtx .getConfirmedBytes ().set (value .getResource ().getSize ());
280
+ }
281
+ sem .release ();
282
+ last = value ;
283
+ }
284
+
285
+ @ Override
286
+ public void onError (Throwable t ) {
287
+ if (t instanceof RuntimeException ) {
288
+ previousError = (RuntimeException ) t ;
289
+ }
290
+ sem .release ();
291
+ }
292
+
293
+ @ Override
294
+ public void onCompleted () {
295
+ if (last != null && last .hasResource ()) {
296
+ resultFuture .set (last );
297
+ }
298
+ sem .release ();
299
+ }
300
+
301
+ void await () {
302
+ try {
303
+ sem .acquire ();
304
+ } catch (InterruptedException e ) {
305
+ if (e .getCause () instanceof RuntimeException ) {
306
+ throw (RuntimeException ) e .getCause ();
307
+ } else {
308
+ throw new RuntimeException (e );
309
+ }
310
+ }
311
+ RuntimeException err = previousError ;
312
+ if (err != null ) {
313
+ previousError = null ;
314
+ throw err ;
315
+ }
316
+ }
317
+ }
165
318
}
0 commit comments