Skip to content

Commit 88ee909

Browse files
OlegDokukarobertroeser
OlegDokuka
authored and
robertroeser
committed
Feature/rate limiter publisher (#672)
* partial rate limiting impl Signed-off-by: Oleh Dokuka <shadowgun@i.ua> * fixes tests Signed-off-by: Oleh Dokuka <shadowgun@i.ua> * Prototyping ratelimited request publisher replacement for coordinated request publisher Signed-off-by: Oleh Dokuka <shadowgun@i.ua>
1 parent 8969401 commit 88ee909

File tree

11 files changed

+412
-23
lines changed

11 files changed

+412
-23
lines changed

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

Lines changed: 11 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@
2727
import io.rsocket.exceptions.Exceptions;
2828
import io.rsocket.frame.*;
2929
import io.rsocket.frame.decoder.PayloadDecoder;
30-
import io.rsocket.internal.LimitableRequestPublisher;
30+
import io.rsocket.internal.RateLimitableRequestPublisher;
3131
import io.rsocket.internal.SynchronizedIntObjectHashMap;
3232
import io.rsocket.internal.UnboundedProcessor;
3333
import io.rsocket.internal.UnicastMonoProcessor;
@@ -47,6 +47,7 @@
4747
import org.reactivestreams.Subscriber;
4848
import org.reactivestreams.Subscription;
4949
import reactor.core.publisher.*;
50+
import reactor.util.concurrent.Queues;
5051

5152
/**
5253
* Requester Side of a RSocket socket. Sends {@link ByteBuf}s to a {@link RSocketResponder} of peer
@@ -60,7 +61,7 @@ class RSocketRequester implements RSocket {
6061
private final PayloadDecoder payloadDecoder;
6162
private final Consumer<Throwable> errorConsumer;
6263
private final StreamIdSupplier streamIdSupplier;
63-
private final IntObjectMap<LimitableRequestPublisher> senders;
64+
private final IntObjectMap<RateLimitableRequestPublisher> senders;
6465
private final IntObjectMap<Processor<Payload, Payload>> receivers;
6566
private final UnboundedProcessor<ByteBuf> sendProcessor;
6667
private final RequesterLeaseHandler leaseHandler;
@@ -131,7 +132,7 @@ private void handleSendProcessorError(Throwable t) {
131132
}
132133
});
133134

134-
senders.values().forEach(LimitableRequestPublisher::cancel);
135+
senders.values().forEach(RateLimitableRequestPublisher::cancel);
135136
}
136137

137138
private void handleSendProcessorCancel(SignalType t) {
@@ -150,7 +151,7 @@ private void handleSendProcessorCancel(SignalType t) {
150151
}
151152
});
152153

153-
senders.values().forEach(LimitableRequestPublisher::cancel);
154+
senders.values().forEach(RateLimitableRequestPublisher::cancel);
154155
}
155156

156157
@Override
@@ -343,8 +344,8 @@ public void accept(long n) {
343344
request
344345
.transform(
345346
f -> {
346-
LimitableRequestPublisher<Payload> wrapped =
347-
LimitableRequestPublisher.wrap(f);
347+
RateLimitableRequestPublisher<Payload> wrapped =
348+
RateLimitableRequestPublisher.wrap(f, Queues.SMALL_BUFFER_SIZE);
348349
// Need to set this to one for first the frame
349350
wrapped.request(1);
350351
senders.put(streamId, wrapped);
@@ -421,7 +422,7 @@ protected void hookOnError(Throwable t) {
421422
.doFinally(
422423
s -> {
423424
receivers.remove(streamId);
424-
LimitableRequestPublisher sender = senders.remove(streamId);
425+
RateLimitableRequestPublisher sender = senders.remove(streamId);
425426
if (sender != null) {
426427
sender.cancel();
427428
}
@@ -489,7 +490,7 @@ private void setTerminationError(Throwable error) {
489490
}
490491

491492
private synchronized void cleanUpLimitableRequestPublisher(
492-
LimitableRequestPublisher<?> limitableRequestPublisher) {
493+
RateLimitableRequestPublisher<?> limitableRequestPublisher) {
493494
try {
494495
limitableRequestPublisher.cancel();
495496
} catch (Throwable t) {
@@ -561,7 +562,7 @@ private void handleFrame(int streamId, FrameType type, ByteBuf frame) {
561562
break;
562563
case CANCEL:
563564
{
564-
LimitableRequestPublisher sender = senders.remove(streamId);
565+
RateLimitableRequestPublisher sender = senders.remove(streamId);
565566
if (sender != null) {
566567
sender.cancel();
567568
}
@@ -572,7 +573,7 @@ private void handleFrame(int streamId, FrameType type, ByteBuf frame) {
572573
break;
573574
case REQUEST_N:
574575
{
575-
LimitableRequestPublisher sender = senders.get(streamId);
576+
RateLimitableRequestPublisher sender = senders.get(streamId);
576577
if (sender != null) {
577578
int n = RequestNFrameFlyweight.requestN(frame);
578579
sender.request(n >= Integer.MAX_VALUE ? Long.MAX_VALUE : n);

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

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
import io.rsocket.exceptions.ApplicationErrorException;
2424
import io.rsocket.frame.*;
2525
import io.rsocket.frame.decoder.PayloadDecoder;
26-
import io.rsocket.internal.LimitableRequestPublisher;
26+
import io.rsocket.internal.RateLimitableRequestPublisher;
2727
import io.rsocket.internal.SynchronizedIntObjectHashMap;
2828
import io.rsocket.internal.UnboundedProcessor;
2929
import io.rsocket.lease.ResponderLeaseHandler;
@@ -35,6 +35,7 @@
3535
import reactor.core.Disposable;
3636
import reactor.core.Exceptions;
3737
import reactor.core.publisher.*;
38+
import reactor.util.concurrent.Queues;
3839

3940
/** Responder side of RSocket. Receives {@link ByteBuf}s from a peer's {@link RSocketRequester} */
4041
class RSocketResponder implements ResponderRSocket {
@@ -46,7 +47,7 @@ class RSocketResponder implements ResponderRSocket {
4647
private final Consumer<Throwable> errorConsumer;
4748
private final ResponderLeaseHandler leaseHandler;
4849

49-
private final IntObjectMap<LimitableRequestPublisher> sendingLimitableSubscriptions;
50+
private final IntObjectMap<RateLimitableRequestPublisher> sendingLimitableSubscriptions;
5051
private final IntObjectMap<Subscription> sendingSubscriptions;
5152
private final IntObjectMap<Processor<Payload, Payload>> channelProcessors;
5253

@@ -435,8 +436,8 @@ private void handleStream(int streamId, Flux<Payload> response, int initialReque
435436
response
436437
.transform(
437438
frameFlux -> {
438-
LimitableRequestPublisher<Payload> payloads =
439-
LimitableRequestPublisher.wrap(frameFlux);
439+
RateLimitableRequestPublisher<Payload> payloads =
440+
RateLimitableRequestPublisher.wrap(frameFlux, Queues.SMALL_BUFFER_SIZE);
440441
sendingLimitableSubscriptions.put(streamId, payloads);
441442
payloads.request(
442443
initialRequestN >= Integer.MAX_VALUE ? Long.MAX_VALUE : initialRequestN);

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

Lines changed: 242 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,242 @@
1+
/*
2+
* Copyright 2015-2018 the original author or authors.
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+
* You may obtain a copy of the License at
7+
*
8+
* http://www.apache.org/licenses/LICENSE-2.0
9+
*
10+
* Unless required by applicable law or agreed to in writing, software
11+
* distributed under the License is distributed on an "AS IS" BASIS,
12+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
* See the License for the specific language governing permissions and
14+
* limitations under the License.
15+
*/
16+
17+
package io.rsocket.internal;
18+
19+
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
20+
import javax.annotation.Nullable;
21+
import org.reactivestreams.Publisher;
22+
import org.reactivestreams.Subscriber;
23+
import org.reactivestreams.Subscription;
24+
import reactor.core.CoreSubscriber;
25+
import reactor.core.publisher.Flux;
26+
import reactor.core.publisher.Operators;
27+
28+
/** */
29+
public class RateLimitableRequestPublisher<T> extends Flux<T> implements Subscription {
30+
31+
private static final int NOT_CANCELED_STATE = 0;
32+
private static final int CANCELED_STATE = 1;
33+
34+
private final Publisher<T> source;
35+
36+
private volatile int canceled;
37+
private static final AtomicIntegerFieldUpdater<RateLimitableRequestPublisher> CANCELED =
38+
AtomicIntegerFieldUpdater.newUpdater(RateLimitableRequestPublisher.class, "canceled");
39+
40+
private final long prefetch;
41+
private final long limit;
42+
43+
private long externalRequested; // need sync
44+
private int pendingToFulfil; // need sync since should be checked/zerroed in onNext
45+
// and increased in request
46+
private int deliveredElements; // no need to sync since increased zerroed only in
47+
// the request method
48+
49+
private boolean subscribed;
50+
51+
private @Nullable Subscription internalSubscription;
52+
53+
private RateLimitableRequestPublisher(Publisher<T> source, long prefetch) {
54+
this.source = source;
55+
this.prefetch = prefetch;
56+
this.limit = prefetch == Integer.MAX_VALUE ? Integer.MAX_VALUE : (prefetch - (prefetch >> 2));
57+
}
58+
59+
public static <T> RateLimitableRequestPublisher<T> wrap(Publisher<T> source, long prefetch) {
60+
return new RateLimitableRequestPublisher<>(source, prefetch);
61+
}
62+
63+
@Override
64+
public void subscribe(CoreSubscriber<? super T> destination) {
65+
synchronized (this) {
66+
if (subscribed) {
67+
throw new IllegalStateException("only one subscriber at a time");
68+
}
69+
70+
subscribed = true;
71+
}
72+
final InnerOperator s = new InnerOperator(destination);
73+
74+
source.subscribe(s);
75+
destination.onSubscribe(s);
76+
}
77+
78+
@Override
79+
public void request(long n) {
80+
synchronized (this) {
81+
long requested = externalRequested;
82+
if (requested == Long.MAX_VALUE) {
83+
return;
84+
}
85+
externalRequested = Operators.addCap(n, requested);
86+
}
87+
88+
requestN();
89+
}
90+
91+
private void requestN() {
92+
final long r;
93+
final Subscription s;
94+
95+
synchronized (this) {
96+
s = internalSubscription;
97+
if (s == null) {
98+
return;
99+
}
100+
101+
final long er = externalRequested;
102+
final long p = prefetch;
103+
final int pendingFulfil = pendingToFulfil;
104+
105+
if (er != Long.MAX_VALUE || p != Integer.MAX_VALUE) {
106+
// shortcut
107+
if (pendingFulfil == p) {
108+
return;
109+
}
110+
111+
r = Math.min(p - pendingFulfil, er);
112+
if (er != Long.MAX_VALUE) {
113+
externalRequested -= r;
114+
}
115+
if (p != Integer.MAX_VALUE) {
116+
pendingToFulfil += r;
117+
}
118+
} else {
119+
r = Long.MAX_VALUE;
120+
}
121+
}
122+
123+
if (r > 0) {
124+
s.request(r);
125+
}
126+
}
127+
128+
public void cancel() {
129+
if (!isCanceled() && CANCELED.compareAndSet(this, NOT_CANCELED_STATE, CANCELED_STATE)) {
130+
Subscription s;
131+
132+
synchronized (this) {
133+
s = internalSubscription;
134+
internalSubscription = null;
135+
subscribed = false;
136+
}
137+
138+
if (s != null) {
139+
s.cancel();
140+
}
141+
}
142+
}
143+
144+
private boolean isCanceled() {
145+
return canceled == CANCELED_STATE;
146+
}
147+
148+
private class InnerOperator implements CoreSubscriber<T>, Subscription {
149+
final Subscriber<? super T> destination;
150+
151+
private InnerOperator(Subscriber<? super T> destination) {
152+
this.destination = destination;
153+
}
154+
155+
@Override
156+
public void onSubscribe(Subscription s) {
157+
synchronized (RateLimitableRequestPublisher.this) {
158+
RateLimitableRequestPublisher.this.internalSubscription = s;
159+
160+
if (isCanceled()) {
161+
s.cancel();
162+
subscribed = false;
163+
RateLimitableRequestPublisher.this.internalSubscription = null;
164+
}
165+
}
166+
167+
requestN();
168+
}
169+
170+
@Override
171+
public void onNext(T t) {
172+
try {
173+
destination.onNext(t);
174+
175+
if (prefetch == Integer.MAX_VALUE) {
176+
return;
177+
}
178+
179+
final long l = limit;
180+
int d = deliveredElements + 1;
181+
182+
if (d == l) {
183+
d = 0;
184+
final long r;
185+
final Subscription s;
186+
187+
synchronized (RateLimitableRequestPublisher.this) {
188+
long er = externalRequested;
189+
s = internalSubscription;
190+
191+
if (s == null) {
192+
return;
193+
}
194+
195+
if (er >= l) {
196+
er -= l;
197+
// keep pendingToFulfil as is since it is eq to prefetch
198+
r = l;
199+
} else {
200+
pendingToFulfil -= l;
201+
if (er > 0) {
202+
r = er;
203+
er = 0;
204+
pendingToFulfil += r;
205+
} else {
206+
r = 0;
207+
}
208+
}
209+
210+
externalRequested = er;
211+
}
212+
213+
if (r > 0) {
214+
s.request(r);
215+
}
216+
}
217+
218+
deliveredElements = d;
219+
} catch (Throwable e) {
220+
onError(e);
221+
}
222+
}
223+
224+
@Override
225+
public void onError(Throwable t) {
226+
destination.onError(t);
227+
}
228+
229+
@Override
230+
public void onComplete() {
231+
destination.onComplete();
232+
}
233+
234+
@Override
235+
public void request(long n) {}
236+
237+
@Override
238+
public void cancel() {
239+
RateLimitableRequestPublisher.this.cancel();
240+
}
241+
}
242+
}

0 commit comments

Comments
 (0)