Skip to content

Commit fe69bda

Browse files
OlegDokukaOlegDokuka
committed
improved Backpressure control
Signed-off-by: Oleh Dokuka <shadowgun@i.ua>
1 parent fce6fe6 commit fe69bda

7 files changed

Lines changed: 100 additions & 102 deletions

File tree

‎rsocket-core/src/main/java/io/rsocket/RSocketClient.java‎

Lines changed: 15 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -326,12 +326,19 @@ public void accept(long n) {
326326
request
327327
.transform(
328328
f -> {
329-
LimitableRequestPublisher<Payload> wrapped =
330-
LimitableRequestPublisher.wrap(
331-
f,
332-
sendProcessor.available() == Long.MAX_VALUE
333-
? Integer.MAX_VALUE
334-
: Queues.SMALL_BUFFER_SIZE);
329+
long available = sendProcessor.available();
330+
LimitableRequestPublisher<Payload> wrapped;
331+
if (available == Long.MAX_VALUE) {
332+
wrapped = LimitableRequestPublisher.wrap(f, Integer.MAX_VALUE);
333+
} else {
334+
long prefetch = Math.max(available, Queues.XS_BUFFER_SIZE);
335+
wrapped =
336+
LimitableRequestPublisher.wrap(f, Queues.XS_BUFFER_SIZE);
337+
if (prefetch > Queues.XS_BUFFER_SIZE) {
338+
long firstRequest = prefetch - Queues.XS_BUFFER_SIZE;
339+
wrapped.internalRequest(firstRequest);
340+
}
341+
}
335342
// Need to set this to one for first the frame
336343
wrapped.request(1);
337344
senders.put(streamId, wrapped);
@@ -418,9 +425,7 @@ protected void hookOnError(Throwable t) {
418425
sendersAsASet.remove(sender);
419426
sender.cancel();
420427
long requested = sender.getInternalRequested();
421-
if (requested > 0) {
422-
shareRequest(requested, sendersAsASet);
423-
}
428+
shareRequest(Math.max(sender.getLimit(), requested), sendersAsASet);
424429
}
425430
});
426431
});
@@ -529,9 +534,7 @@ private void handleFrame(int streamId, FrameType type, ByteBuf frame) {
529534
sendersAsASet.remove(sender);
530535
sender.cancel();
531536
long requested = sender.getInternalRequested();
532-
if (requested > 0) {
533-
shareRequest(requested, sendersAsASet);
534-
}
537+
shareRequest(Math.max(sender.getLimit(), requested), sendersAsASet);
535538
}
536539
break;
537540
}

‎rsocket-core/src/main/java/io/rsocket/RSocketServer.java‎

Lines changed: 22 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -440,17 +440,23 @@ private void handleStream(int streamId, Flux<Payload> response, int initialReque
440440
response
441441
.transform(
442442
frameFlux -> {
443-
LimitableRequestPublisher<Payload> payloads =
444-
LimitableRequestPublisher.wrap(
445-
frameFlux,
446-
sendProcessor.available() == Long.MAX_VALUE
447-
? Integer.MAX_VALUE
448-
: Queues.SMALL_BUFFER_SIZE);
449-
sendingSubscriptions.put(streamId, payloads);
450-
sendingLimitableSubscriptions.add(payloads);
451-
payloads.request(
443+
long available = sendProcessor.available();
444+
LimitableRequestPublisher<Payload> wrapped;
445+
if (available == Long.MAX_VALUE) {
446+
wrapped = LimitableRequestPublisher.wrap(frameFlux, Integer.MAX_VALUE);
447+
} else {
448+
long prefetch = Math.max(available, Queues.XS_BUFFER_SIZE);
449+
wrapped = LimitableRequestPublisher.wrap(frameFlux, Queues.XS_BUFFER_SIZE);
450+
if (prefetch > Queues.XS_BUFFER_SIZE) {
451+
long firstRequest = prefetch - Queues.XS_BUFFER_SIZE;
452+
wrapped.internalRequest(firstRequest);
453+
}
454+
}
455+
sendingSubscriptions.put(streamId, wrapped);
456+
sendingLimitableSubscriptions.add(wrapped);
457+
wrapped.request(
452458
initialRequestN >= Integer.MAX_VALUE ? Long.MAX_VALUE : initialRequestN);
453-
return payloads;
459+
return wrapped;
454460
})
455461
.subscribe(
456462
new BaseSubscriber<Payload>() {
@@ -486,9 +492,8 @@ protected void hookFinally(SignalType type) {
486492
(LimitableRequestPublisher) sendingSubscriptions.remove(streamId);
487493
sendingLimitableSubscriptions.remove(subscription);
488494
long requested = subscription.getInternalRequested();
489-
if (requested > 0) {
490-
shareRequest(requested, sendingLimitableSubscriptions);
491-
}
495+
shareRequest(
496+
Math.max(subscription.getLimit(), requested), sendingLimitableSubscriptions);
492497
}
493498
});
494499
}
@@ -533,11 +538,10 @@ private void handleCancelFrame(int streamId) {
533538
if (subscription instanceof LimitableRequestPublisher) {
534539
sendingLimitableSubscriptions.remove(subscription);
535540

536-
long requested =
537-
((LimitableRequestPublisher) subscription).getInternalRequested();
538-
if (requested > 0) {
539-
shareRequest(requested, sendingLimitableSubscriptions);
540-
}
541+
long requested = ((LimitableRequestPublisher) subscription).getInternalRequested();
542+
shareRequest(
543+
Math.max(((LimitableRequestPublisher) subscription).getLimit(), requested),
544+
sendingLimitableSubscriptions);
541545
}
542546
}
543547
}

‎rsocket-core/src/main/java/io/rsocket/internal/LimitableRequestPublisher.java‎

Lines changed: 49 additions & 59 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
package io.rsocket.internal;
1818

1919
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
20+
import java.util.concurrent.atomic.AtomicLongFieldUpdater;
2021
import javax.annotation.Nullable;
2122
import org.reactivestreams.Publisher;
2223
import org.reactivestreams.Subscriber;
@@ -42,8 +43,17 @@ public class LimitableRequestPublisher<T> extends Flux<T> implements Subscriptio
4243

4344
private int produced;
4445

45-
private long internalRequested;
46-
private long externalRequested;
46+
private volatile long internalRequested;
47+
private static final AtomicLongFieldUpdater<LimitableRequestPublisher> INTERNAL_REQUESTED =
48+
AtomicLongFieldUpdater.newUpdater(LimitableRequestPublisher.class, "internalRequested");
49+
50+
private volatile long externalRequested;
51+
private static final AtomicLongFieldUpdater<LimitableRequestPublisher> EXTERNAL_REQUESTED =
52+
AtomicLongFieldUpdater.newUpdater(LimitableRequestPublisher.class, "externalRequested");
53+
54+
private volatile int wip;
55+
private static final AtomicIntegerFieldUpdater<LimitableRequestPublisher> WIP =
56+
AtomicIntegerFieldUpdater.newUpdater(LimitableRequestPublisher.class, "wip");
4757

4858
private boolean subscribed;
4959

@@ -63,10 +73,6 @@ public int getLimit() {
6373
return limit;
6474
}
6575

66-
public int getProduced() {
67-
return produced;
68-
}
69-
7076
public long getInternalRequested() {
7177
return internalRequested;
7278
}
@@ -97,73 +103,57 @@ public int getPrefetch() {
97103
}
98104

99105
public void internalRequest(long n) {
100-
int p;
101-
long r;
102-
103-
synchronized (this) {
104-
p = produced;
105-
long requested = internalRequested;
106-
if (requested == Long.MAX_VALUE) {
107-
return;
108-
}
109-
r = Operators.addCap(n, requested);
110-
internalRequested = r;
111-
}
112-
113-
if (r >= limit) {
114-
requestN();
115-
}
106+
Operators.addCap(INTERNAL_REQUESTED, this, n);
107+
requestN();
116108
}
117109

118110
@Override
119111
public void request(long n) {
120-
synchronized (this) {
121-
long requested = externalRequested;
122-
if (requested == Long.MAX_VALUE) {
123-
return;
124-
}
125-
externalRequested = Operators.addCap(n, requested);
126-
}
127-
112+
// if (Operators.addCap(EXTERNAL_REQUESTED, this, n) == 0) {
113+
Operators.addCap(EXTERNAL_REQUESTED, this, n);
128114
requestN();
115+
// }
129116
}
130117

131118
private void requestN() {
132-
long r;
133-
final Subscription s;
134-
135-
synchronized (this) {
136-
s = internalSubscription;
137-
if (s == null) {
138-
return;
139-
}
140-
141-
long er = externalRequested;
142-
long ir = internalRequested;
143-
144-
if (er != Long.MAX_VALUE || ir != Long.MAX_VALUE) {
145-
int limit = this.limit;
146-
r = Math.min(ir, limit);
119+
if (WIP.getAndIncrement(this) > 0) {
120+
return;
121+
}
147122

148-
if (r != limit) {
149-
return;
150-
}
123+
long r;
124+
final Subscription s = internalSubscription;
125+
int missed = 1;
151126

152-
r = Math.min(r, er);
127+
long er = externalRequested;
128+
long ir = internalRequested;
129+
int limit = this.limit;
153130

154-
if (er != Long.MAX_VALUE) {
155-
externalRequested -= r;
156-
}
157-
if (ir != Long.MAX_VALUE) {
158-
internalRequested -= r;
131+
while (true) {
132+
if (s != null) {
133+
if (er != Long.MAX_VALUE || ir != Long.MAX_VALUE) {
134+
if (ir > limit) {
135+
r = Math.min(ir, er);
136+
137+
if (er != Long.MAX_VALUE) {
138+
er = EXTERNAL_REQUESTED.addAndGet(this, -r);
139+
}
140+
if (ir != Long.MAX_VALUE) {
141+
ir = INTERNAL_REQUESTED.addAndGet(this, -r);
142+
}
143+
144+
if (r > 0) {
145+
s.request(r);
146+
}
147+
}
148+
} else {
149+
s.request(Long.MAX_VALUE);
159150
}
160-
} else {
161-
r = Long.MAX_VALUE;
162151
}
163-
}
164152

165-
if (r > 0) {
166-
s.request(r);
153+
missed = WIP.addAndGet(this, -missed);
154+
if (missed == 0) {
155+
break;
156+
}
167157
}
168158
}
169159

‎rsocket-core/src/main/java/io/rsocket/util/BackpressureUtils.java‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,7 @@ public static void shareRequest(
3131
int count = 0;
3232
boolean any = false;
3333

34-
long minimumPrefetch = requested > length ? requested / length : 1;
34+
// long minimumPrefetch = (requested << 8) > length ? requested / length : 1;
3535

3636
while (requested > 0) {
3737
LimitableRequestPublisher subscription = values.get(i);
@@ -42,16 +42,19 @@ public static void shareRequest(
4242
&& subscription.getInternalRequested() == 0)
4343
|| count >= length) {
4444
int prefetch = subscription.getLimit();
45-
subscription.internalRequest(minimumPrefetch > prefetch ? minimumPrefetch : prefetch);
45+
subscription.internalRequest(prefetch);
4646
requested -= prefetch;
4747
}
4848
}
4949

5050
count++;
5151
i = ++i % length;
5252

53-
if (count >= length && !any) {
54-
return;
53+
if (count % length == 0) {
54+
if (!any) {
55+
return;
56+
}
57+
any = false;
5558
}
5659
}
5760
}

‎rsocket-core/src/test/java/io/rsocket/internal/LimitableRequestPublisherTest.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ class LimitableRequestPublisherTest {
1515
public void requestLimitRacingTest() throws InterruptedException {
1616
Queue<Long> requests = new ArrayDeque<>(10000);
1717
LimitableRequestPublisher<Object> limitableRequestPublisher =
18-
LimitableRequestPublisher.wrap(DirectProcessor.create().doOnRequest(requests::add), 32);
18+
LimitableRequestPublisher.wrap(DirectProcessor.create().doOnRequest(requests::add), 4);
1919

2020
Runnable request1 = () -> limitableRequestPublisher.request(1);
2121
Runnable request2 = () -> limitableRequestPublisher.internalRequest(2);

‎rsocket-examples/src/test/java/io/rsocket/integration/TestingStreaming.java‎

Lines changed: 5 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -24,7 +24,6 @@
2424
import io.rsocket.transport.local.LocalServerTransport;
2525
import io.rsocket.util.DefaultPayload;
2626
import java.time.Duration;
27-
import java.util.concurrent.atomic.AtomicInteger;
2827
import java.util.function.Supplier;
2928
import org.junit.Test;
3029
import reactor.core.publisher.Flux;
@@ -124,12 +123,11 @@ private Flux<Payload> consumer(String s) {
124123
.transport(clientSupplier)
125124
.start()
126125
.flatMapMany(
127-
rSocket -> {
128-
AtomicInteger count = new AtomicInteger();
129-
return Flux.range(1, 100)
130-
.flatMap(
131-
i -> rSocket.requestStream(DefaultPayload.create("i -> " + i)).take(100), 1);
132-
});
126+
rSocket ->
127+
Flux.range(1, 100)
128+
.flatMap(
129+
i -> rSocket.requestStream(DefaultPayload.create("i -> " + i)).take(100),
130+
1));
133131
}
134132

135133
@Test

‎rsocket-transport-netty/src/test/java/io/rsocket/transport/netty/TcpPing.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -89,7 +89,7 @@ private static PingClient newPingClient(boolean isResumable) {
8989
clientRSocketFactory
9090
.frameDecoder(PayloadDecoder.ZERO_COPY)
9191
.keepAlive(Duration.ofMinutes(1), Duration.ofMinutes(30), 3)
92-
.transport(TcpClientTransport.create(port))
92+
.transport(TcpClientTransport.create("localhost", port))
9393
.start();
9494

9595
return new PingClient(rSocket);

0 commit comments

Comments
 (0)