Fix single upload retry. (#12834)

Azure · Jul 6, 2020 · b730543 · b730543
1 parent 714c26b
commit b730543
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/...storage-common/src/main/java/com/azure/storage/common/implementation/PayloadSizeGate.java b/...storage-common/src/main/java/com/azure/storage/common/implementation/PayloadSizeGate.java
@@ -66,7 +66,8 @@ Flux<ByteBuffer> write(ByteBuffer buf) {
      */
     Flux<ByteBuffer> flush() {
         if (byteBuffers != null) {
-            Flux<ByteBuffer> result = dequeuingFlux(byteBuffers);
+            // We return Flux from iterable in this case to support retries on single upload.
+            Flux<ByteBuffer> result = Flux.fromIterable(byteBuffers);
             byteBuffers = null;
             return result;
         } else {

diff --git a/...ure-storage-common/src/main/java/com/azure/storage/common/implementation/UploadUtils.java b/...ure-storage-common/src/main/java/com/azure/storage/common/implementation/UploadUtils.java
@@ -44,14 +44,18 @@ public static <T> Mono<Response<T>> uploadFullOrChunked(final Flux<ByteBuffer> d
             .filter(ByteBuffer::hasRemaining)
             // The gate buffers data until threshold is breached.
             .concatMap(gate::write)
-            .concatWith(Flux.defer(gate::flush))
             // First buffer is emitted after threshold is breached or there's no more data.
             // Therefore we can make a decision how to upload data on first element.
             .switchOnFirst((signal, flux) -> {
                 if (gate.isThresholdBreached()) {
-                    return uploadInChunks.apply(flux);
+                    // In this case we can pass a flux that can have just one subscriber because
+                    // the chunked upload is going to cache the data downstream before sending chunks over the wire.
+                    return uploadInChunks.apply(flux.concatWith(Flux.defer(gate::flush)));
                 } else {
-                    return uploadFull.apply(flux, gate.size());
+                    // In this case gate contains all the data cached.
+                    // The flux passed to this lambda allows only one subscriber. Therefore we substitute it
+                    // with flux coming from gate which is based of iterable and can be subscribed again.
+                    return uploadFull.apply(gate.flush(), gate.size());
                 }
             })
             .next()