Skip to content

Commit f8c1cb0

Browse files
zoewanggzoewangg
authored
Add AddingTrailingDataSubscriber to allow users to send additional data t… (#4366)
* Add AdditionalDataSubscriber to allow users to send additional data to the downstream subscriber * Support iterable
1 parent 34aa46d commit f8c1cb0

File tree

5 files changed

+377
-0
lines changed

5 files changed

+377
-0
lines changed

‎core/sdk-core/src/main/java/software/amazon/awssdk/core/async/SdkPublisher.java

+14
Original file line numberDiff line numberDiff line change
@@ -20,10 +20,12 @@
2020
import java.util.function.Consumer;
2121
import java.util.function.Function;
2222
import java.util.function.Predicate;
23+
import java.util.function.Supplier;
2324
import org.reactivestreams.Publisher;
2425
import org.reactivestreams.Subscriber;
2526
import org.reactivestreams.Subscription;
2627
import software.amazon.awssdk.annotations.SdkPublicApi;
28+
import software.amazon.awssdk.utils.async.AddingTrailingDataSubscriber;
2729
import software.amazon.awssdk.utils.async.BufferingSubscriber;
2830
import software.amazon.awssdk.utils.async.EventListeningSubscriber;
2931
import software.amazon.awssdk.utils.async.FilteringSubscriber;
@@ -118,6 +120,18 @@ default SdkPublisher<T> limit(int limit) {
118120
return subscriber -> subscribe(new LimitingSubscriber<>(subscriber, limit));
119121
}
120122

123+
124+
/**
125+
* Creates a new publisher that emits trailing events provided by {@code trailingDataSupplier} in addition to the
126+
* published events.
127+
*
128+
* @param trailingDataSupplier supplier to provide the trailing data
129+
* @return New publisher that will publish additional events
130+
*/
131+
default SdkPublisher<T> addTrailingData(Supplier<Iterable<T>> trailingDataSupplier) {
132+
return subscriber -> subscribe(new AddingTrailingDataSubscriber<T>(subscriber, trailingDataSupplier));
133+
}
134+
121135
/**
122136
* Add a callback that will be invoked after this publisher invokes {@link Subscriber#onComplete()}.
123137
*

‎core/sdk-core/src/test/java/software/amazon/awssdk/core/async/SdkPublishersTest.java

+18
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
import java.nio.ByteBuffer;
2222
import java.nio.charset.StandardCharsets;
2323
import java.util.ArrayList;
24+
import java.util.Arrays;
2425
import java.util.List;
2526
import java.util.concurrent.CompletableFuture;
2627
import java.util.concurrent.ExecutionException;
@@ -141,6 +142,23 @@ public void flatMapIterableHandlesError() {
141142
.hasCause(exception);
142143
}
143144

145+
@Test
146+
public void addTrailingData_handlesCorrectly() {
147+
FakeSdkPublisher<String> fakePublisher = new FakeSdkPublisher<>();
148+
149+
FakeStringSubscriber fakeSubscriber = new FakeStringSubscriber();
150+
fakePublisher.addTrailingData(() -> Arrays.asList("two", "three"))
151+
.subscribe(fakeSubscriber);
152+
153+
fakePublisher.publish("one");
154+
fakePublisher.complete();
155+
156+
assertThat(fakeSubscriber.recordedEvents()).containsExactly("one", "two", "three");
157+
assertThat(fakeSubscriber.isComplete()).isTrue();
158+
assertThat(fakeSubscriber.isError()).isFalse();
159+
}
160+
161+
144162
private final static class FakeByteBufferSubscriber implements Subscriber<ByteBuffer> {
145163
private final List<String> recordedEvents = new ArrayList<>();
146164

‎utils/src/main/java/software/amazon/awssdk/utils/async/AddingTrailingDataSubscriber.java

+171
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,171 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.utils.async;
17+
18+
import java.util.Iterator;
19+
import java.util.concurrent.atomic.AtomicLong;
20+
import java.util.function.Supplier;
21+
import org.reactivestreams.Subscriber;
22+
import org.reactivestreams.Subscription;
23+
import software.amazon.awssdk.annotations.SdkProtectedApi;
24+
import software.amazon.awssdk.utils.Logger;
25+
import software.amazon.awssdk.utils.Validate;
26+
27+
/**
28+
* Allows to send trailing data before invoking onComplete on the downstream subscriber.
29+
* trailingDataIterable will be created when the upstream subscriber has called onComplete.
30+
*/
31+
@SdkProtectedApi
32+
public class AddingTrailingDataSubscriber<T> extends DelegatingSubscriber<T, T> {
33+
private static final Logger log = Logger.loggerFor(AddingTrailingDataSubscriber.class);
34+
35+
/**
36+
* The subscription to the upstream subscriber.
37+
*/
38+
private Subscription upstreamSubscription;
39+
40+
/**
41+
* The amount of unfulfilled demand the downstream subscriber has opened against us.
42+
*/
43+
private final AtomicLong downstreamDemand = new AtomicLong(0);
44+
45+
/**
46+
* Whether the upstream subscriber has called onComplete on us.
47+
*/
48+
private volatile boolean onCompleteCalledByUpstream = false;
49+
50+
/**
51+
* Whether the upstream subscriber has called onError on us.
52+
*/
53+
private volatile boolean onErrorCalledByUpstream = false;
54+
55+
/**
56+
* Whether we have called onComplete on the downstream subscriber.
57+
*/
58+
private volatile boolean onCompleteCalledOnDownstream = false;
59+
60+
private final Supplier<Iterable<T>> trailingDataIterableSupplier;
61+
private Iterator<T> trailingDataIterator;
62+
63+
public AddingTrailingDataSubscriber(Subscriber<? super T> subscriber,
64+
Supplier<Iterable<T>> trailingDataIterableSupplier) {
65+
super(Validate.paramNotNull(subscriber, "subscriber"));
66+
this.trailingDataIterableSupplier = Validate.paramNotNull(trailingDataIterableSupplier, "trailingDataIterableSupplier");
67+
}
68+
69+
@Override
70+
public void onSubscribe(Subscription subscription) {
71+
72+
if (upstreamSubscription != null) {
73+
log.warn(() -> "Received duplicate subscription, cancelling the duplicate.", new IllegalStateException());
74+
subscription.cancel();
75+
return;
76+
}
77+
78+
upstreamSubscription = subscription;
79+
80+
subscriber.onSubscribe(new Subscription() {
81+
82+
@Override
83+
public void request(long l) {
84+
if (onErrorCalledByUpstream || onCompleteCalledOnDownstream) {
85+
return;
86+
}
87+
88+
addDownstreamDemand(l);
89+
90+
if (onCompleteCalledByUpstream) {
91+
sendTrailingDataAndCompleteIfNeeded();
92+
return;
93+
}
94+
upstreamSubscription.request(l);
95+
}
96+
97+
@Override
98+
public void cancel() {
99+
upstreamSubscription.cancel();
100+
}
101+
});
102+
}
103+
104+
@Override
105+
public void onError(Throwable throwable) {
106+
onErrorCalledByUpstream = true;
107+
subscriber.onError(throwable);
108+
}
109+
110+
@Override
111+
public void onNext(T t) {
112+
Validate.paramNotNull(t, "item");
113+
downstreamDemand.decrementAndGet();
114+
subscriber.onNext(t);
115+
}
116+
117+
@Override
118+
public void onComplete() {
119+
onCompleteCalledByUpstream = true;
120+
sendTrailingDataAndCompleteIfNeeded();
121+
}
122+
123+
private void addDownstreamDemand(long l) {
124+
125+
if (l > 0) {
126+
downstreamDemand.getAndUpdate(current -> {
127+
long newValue = current + l;
128+
return newValue >= 0 ? newValue : Long.MAX_VALUE;
129+
});
130+
} else {
131+
upstreamSubscription.cancel();
132+
onError(new IllegalArgumentException("Demand must not be negative"));
133+
}
134+
}
135+
136+
private synchronized void sendTrailingDataAndCompleteIfNeeded() {
137+
if (onCompleteCalledOnDownstream) {
138+
return;
139+
}
140+
141+
if (trailingDataIterator == null) {
142+
Iterable<T> supplier = trailingDataIterableSupplier.get();
143+
if (supplier == null) {
144+
completeDownstreamSubscriber();
145+
return;
146+
}
147+
148+
trailingDataIterator = supplier.iterator();
149+
}
150+
151+
sendTrailingDataIfNeeded();
152+
153+
if (!trailingDataIterator.hasNext()) {
154+
completeDownstreamSubscriber();
155+
}
156+
}
157+
158+
private void sendTrailingDataIfNeeded() {
159+
long demand = downstreamDemand.get();
160+
161+
while (trailingDataIterator.hasNext() && demand > 0) {
162+
subscriber.onNext(trailingDataIterator.next());
163+
demand = downstreamDemand.decrementAndGet();
164+
}
165+
}
166+
167+
private void completeDownstreamSubscriber() {
168+
subscriber.onComplete();
169+
onCompleteCalledOnDownstream = true;
170+
}
171+
}

‎utils/src/test/java/software/amazon/awssdk/utils/async/AddingTrailingDataSubscriberTckTest.java

+75
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,75 @@
1+
/*
2+
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.awssdk.utils.async;
17+
18+
import java.util.Arrays;
19+
import java.util.concurrent.CompletableFuture;
20+
import org.reactivestreams.Subscriber;
21+
import org.reactivestreams.Subscription;
22+
import org.reactivestreams.tck.SubscriberWhiteboxVerification;
23+
import org.reactivestreams.tck.TestEnvironment;
24+
25+
public class AddingTrailingDataSubscriberTckTest extends SubscriberWhiteboxVerification<Integer> {
26+
protected AddingTrailingDataSubscriberTckTest() {
27+
super(new TestEnvironment());
28+
}
29+
30+
@Override
31+
public Subscriber<Integer> createSubscriber(WhiteboxSubscriberProbe<Integer> probe) {
32+
Subscriber<Integer> foo = new SequentialSubscriber<>(s -> {}, new CompletableFuture<>());
33+
34+
return new AddingTrailingDataSubscriber<Integer>(foo, () -> Arrays.asList(0, 1, 2)) {
35+
@Override
36+
public void onError(Throwable throwable) {
37+
super.onError(throwable);
38+
probe.registerOnError(throwable);
39+
}
40+
41+
@Override
42+
public void onSubscribe(Subscription subscription) {
43+
super.onSubscribe(subscription);
44+
probe.registerOnSubscribe(new SubscriberPuppet() {
45+
@Override
46+
public void triggerRequest(long elements) {
47+
subscription.request(elements);
48+
}
49+
50+
@Override
51+
public void signalCancel() {
52+
subscription.cancel();
53+
}
54+
});
55+
}
56+
57+
@Override
58+
public void onNext(Integer nextItem) {
59+
super.onNext(nextItem);
60+
probe.registerOnNext(nextItem);
61+
}
62+
63+
@Override
64+
public void onComplete() {
65+
super.onComplete();
66+
probe.registerOnComplete();
67+
}
68+
};
69+
}
70+
71+
@Override
72+
public Integer createElement(int i) {
73+
return i;
74+
}
75+
}

0 commit comments

Comments
 (0)