Skip to content

Commit 1a0522b

Browse files
committed
DataBufferUtils does not release DataBuffer on error cases
This commit makes sure that in DataBufferUtils.write, any received data buffers are returned as part of the returned flux, even when an error occurs or is received. Issue: SPR-16782
1 parent 1da4d50 commit 1a0522b

File tree

2 files changed

+240
-25
lines changed

2 files changed

+240
-25
lines changed

spring-core/src/main/java/org/springframework/core/io/buffer/DataBufferUtils.java

Lines changed: 48 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -31,6 +31,7 @@
3131
import java.util.concurrent.Callable;
3232
import java.util.concurrent.atomic.AtomicBoolean;
3333
import java.util.concurrent.atomic.AtomicLong;
34+
import java.util.concurrent.atomic.AtomicReference;
3435
import java.util.function.Consumer;
3536
import java.util.function.IntPredicate;
3637

@@ -336,6 +337,7 @@ public static Flux<DataBuffer> write(Publisher<DataBuffer> source, WritableByteC
336337
sink.next(dataBuffer);
337338
}
338339
catch (IOException ex) {
340+
sink.next(dataBuffer);
339341
sink.error(ex);
340342
}
341343

@@ -355,6 +357,26 @@ public static Flux<DataBuffer> write(Publisher<DataBuffer> source, WritableByteC
355357
* @param channel the channel to write to
356358
* @return a flux containing the same buffers as in {@code source}, that starts the writing
357359
* process when subscribed to, and that publishes any writing errors and the completion signal
360+
* @since 5.1
361+
*/
362+
public static Flux<DataBuffer> write(
363+
Publisher<DataBuffer> source, AsynchronousFileChannel channel) {
364+
return write(source, channel, 0);
365+
}
366+
367+
368+
/**
369+
* Write the given stream of {@link DataBuffer DataBuffers} to the given {@code AsynchronousFileChannel}.
370+
* Does <strong>not</strong> close the channel when the flux is terminated, and does
371+
* <strong>not</strong> {@linkplain #release(DataBuffer) release} the data buffers in the
372+
* source. If releasing is required, then subscribe to the returned {@code Flux} with a
373+
* {@link #releaseConsumer()}.
374+
* <p>Note that the writing process does not start until the returned {@code Flux} is subscribed to.
375+
* @param source the stream of data buffers to be written
376+
* @param channel the channel to write to
377+
* @param position the file position at which the write is to begin; must be non-negative
378+
* @return a flux containing the same buffers as in {@code source}, that starts the writing
379+
* process when subscribed to, and that publishes any writing errors and the completion signal
358380
*/
359381
public static Flux<DataBuffer> write(
360382
Publisher<DataBuffer> source, AsynchronousFileChannel channel, long position) {
@@ -610,10 +632,11 @@ private static class AsynchronousFileChannelWriteCompletionHandler extends BaseS
610632

611633
private final AtomicBoolean completed = new AtomicBoolean();
612634

635+
private final AtomicReference<Throwable> error = new AtomicReference<>();
636+
613637
private final AtomicLong position;
614638

615-
@Nullable
616-
private DataBuffer dataBuffer;
639+
private final AtomicReference<DataBuffer> dataBuffer = new AtomicReference<>();
617640

618641
public AsynchronousFileChannelWriteCompletionHandler(
619642
FluxSink<DataBuffer> sink, AsynchronousFileChannel channel, long position) {
@@ -630,21 +653,27 @@ protected void hookOnSubscribe(Subscription subscription) {
630653

631654
@Override
632655
protected void hookOnNext(DataBuffer value) {
633-
this.dataBuffer = value;
656+
if (!this.dataBuffer.compareAndSet(null, value)) {
657+
throw new IllegalStateException();
658+
}
634659
ByteBuffer byteBuffer = value.asByteBuffer();
635660
this.channel.write(byteBuffer, this.position.get(), byteBuffer, this);
636661
}
637662

638663
@Override
639664
protected void hookOnError(Throwable throwable) {
640-
this.sink.error(throwable);
665+
this.error.set(throwable);
666+
667+
if (this.dataBuffer.get() == null) {
668+
this.sink.error(throwable);
669+
}
641670
}
642671

643672
@Override
644673
protected void hookOnComplete() {
645674
this.completed.set(true);
646675

647-
if (this.dataBuffer == null) {
676+
if (this.dataBuffer.get() == null) {
648677
this.sink.complete();
649678
}
650679
}
@@ -656,11 +685,13 @@ public void completed(Integer written, ByteBuffer byteBuffer) {
656685
this.channel.write(byteBuffer, pos, byteBuffer, this);
657686
return;
658687
}
659-
if (this.dataBuffer != null) {
660-
this.sink.next(this.dataBuffer);
661-
this.dataBuffer = null;
688+
sinkDataBuffer();
689+
690+
Throwable throwable = this.error.get();
691+
if (throwable != null) {
692+
this.sink.error(throwable);
662693
}
663-
if (this.completed.get()) {
694+
else if (this.completed.get()) {
664695
this.sink.complete();
665696
}
666697
else {
@@ -670,8 +701,16 @@ public void completed(Integer written, ByteBuffer byteBuffer) {
670701

671702
@Override
672703
public void failed(Throwable exc, ByteBuffer byteBuffer) {
704+
sinkDataBuffer();
673705
this.sink.error(exc);
674706
}
707+
708+
private void sinkDataBuffer() {
709+
DataBuffer dataBuffer = this.dataBuffer.get();
710+
Assert.state(dataBuffer != null, "DataBuffer should not be null");
711+
this.sink.next(dataBuffer);
712+
this.dataBuffer.set(null);
713+
}
675714
}
676715

677716
/**

0 commit comments

Comments
 (0)