Skip to content
This repository was archived by the owner on Jun 21, 2023. It is now read-only.

Commit 392b633

Browse files
segabrielsegabriel
authored
Merge pull request #102 from /issues/101/gateway-ws-interceptor
#101 enhance ws gateway inerceptors
2 parents 135f1a4 + 5ec455e commit 392b633

File tree

5 files changed

+178
-118
lines changed

5 files changed

+178
-118
lines changed

‎services-gateway-netty/src/main/java/io/scalecube/services/gateway/GatewaySessionHandler.java‎

Lines changed: 43 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,11 @@
11
package io.scalecube.services.gateway;
22

3+
import io.netty.buffer.ByteBuf;
34
import io.scalecube.services.api.ServiceMessage;
45
import io.scalecube.services.gateway.ws.GatewayMessage;
56
import org.slf4j.Logger;
67
import org.slf4j.LoggerFactory;
8+
import reactor.util.context.Context;
79

810
public interface GatewaySessionHandler<M> {
911

@@ -25,30 +27,63 @@ default M mapMessage(GatewaySession session, M req) {
2527
return req;
2628
}
2729

30+
/**
31+
* Request mapper function.
32+
*
33+
* @param session session
34+
* @param byteBuf request buffer
35+
* @param context subscriber context
36+
* @return context
37+
*/
38+
default Context onRequest(GatewaySession session, ByteBuf byteBuf, Context context) {
39+
return context;
40+
}
41+
42+
/**
43+
* On response handler.
44+
*
45+
* @param session session
46+
* @param byteBuf response buffer
47+
* @param message response message
48+
* @param context subscriber context
49+
*/
50+
default void onResponse(
51+
GatewaySession session, ByteBuf byteBuf, GatewayMessage message, Context context) {
52+
// no-op
53+
}
54+
2855
/**
2956
* Error handler function.
3057
*
3158
* @param session webscoket session (not null)
3259
* @param throwable an exception that occurred (not null)
33-
* @param req request message (optional)
34-
* @param resp response message (optional)
60+
* @param context subscriber context
3561
*/
36-
default void onError(GatewaySession session, Throwable throwable, M req, M resp) {
62+
default void onError(GatewaySession session, Throwable throwable, Context context) {
3763
LOGGER.error(
38-
"Exception occurred on session={}, on request: {}, on response: {}, cause:",
64+
"Exception occurred on session: {}, on context: {}, cause:",
3965
session.sessionId(),
40-
req,
41-
resp,
66+
context,
4267
throwable);
4368
}
4469

70+
/**
71+
* On session error.
72+
*
73+
* @param session webscoket session (not null)
74+
* @param throwable an exception that occurred (not null)
75+
*/
76+
default void onSessionError(GatewaySession session, Throwable throwable) {
77+
LOGGER.error("Exception occurred on session: {}, cause:", session.sessionId(), throwable);
78+
}
79+
4580
/**
4681
* On session open handler.
4782
*
4883
* @param session websocket session (not null)
4984
*/
5085
default void onSessionOpen(GatewaySession session) {
51-
LOGGER.info("Session opened: " + session);
86+
LOGGER.info("Session opened: {}", session);
5287
}
5388

5489
/**
@@ -57,6 +92,6 @@ default void onSessionOpen(GatewaySession session) {
5792
* @param session websocket session (not null)
5893
*/
5994
default void onSessionClose(GatewaySession session) {
60-
LOGGER.info("Session closed: " + session);
95+
LOGGER.info("Session closed: {}", session);
6196
}
6297
}

‎services-gateway-netty/src/main/java/io/scalecube/services/gateway/ws/WebsocketContextException.java‎

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
package io.scalecube.services.gateway.ws;
2+
3+
import io.scalecube.services.gateway.ReferenceCountUtil;
4+
import java.util.Optional;
5+
6+
public class WebsocketContextException extends RuntimeException {
7+
8+
private final GatewayMessage request;
9+
private final GatewayMessage response;
10+
11+
private WebsocketContextException(
12+
Throwable cause, GatewayMessage request, GatewayMessage response) {
13+
super(cause);
14+
this.request = request;
15+
this.response = response;
16+
}
17+
18+
public static WebsocketContextException badRequest(String errorMessage, GatewayMessage request) {
19+
return new WebsocketContextException(
20+
new io.scalecube.services.exceptions.BadRequestException(errorMessage), request, null);
21+
}
22+
23+
public static WebsocketContextException wrap(
24+
Throwable th, GatewayMessage request, GatewayMessage response) {
25+
return new WebsocketContextException(th, request, response);
26+
}
27+
28+
public GatewayMessage request() {
29+
return request;
30+
}
31+
32+
public GatewayMessage response() {
33+
return response;
34+
}
35+
36+
/**
37+
* Releases request data if any.
38+
*
39+
* @return self
40+
*/
41+
public WebsocketContextException releaseRequest() {
42+
Optional.ofNullable(request)
43+
.map(GatewayMessage::data)
44+
.ifPresent(ReferenceCountUtil::safestRelease);
45+
return this;
46+
}
47+
}

‎services-gateway-netty/src/main/java/io/scalecube/services/gateway/ws/WebsocketGatewayAcceptor.java‎

Lines changed: 71 additions & 67 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
package io.scalecube.services.gateway.ws;
22

3+
import io.netty.buffer.ByteBuf;
34
import io.scalecube.services.ServiceCall;
45
import io.scalecube.services.api.ServiceMessage;
56
import io.scalecube.services.exceptions.DefaultErrorMapper;
@@ -19,7 +20,7 @@
1920
import reactor.netty.http.server.HttpServerResponse;
2021
import reactor.netty.http.websocket.WebsocketInbound;
2122
import reactor.netty.http.websocket.WebsocketOutbound;
22-
23+
import reactor.util.context.Context;
2324

2425
public class WebsocketGatewayAcceptor
2526
implements BiFunction<HttpServerRequest, HttpServerResponse, Publisher<Void>> {
@@ -48,41 +49,55 @@ public WebsocketGatewayAcceptor(
4849
public Publisher<Void> apply(HttpServerRequest httpRequest, HttpServerResponse httpResponse) {
4950
return httpResponse.sendWebsocket(
5051
(WebsocketInbound inbound, WebsocketOutbound outbound) ->
51-
onConnect(new WebsocketGatewaySession(messageCodec, httpRequest, inbound, outbound)));
52+
onConnect(
53+
new WebsocketGatewaySession(
54+
messageCodec, httpRequest, inbound, outbound, gatewayHandler)));
5255
}
5356

5457
private Mono<Void> onConnect(WebsocketGatewaySession session) {
5558
gatewayHandler.onSessionOpen(session);
5659

5760
session
5861
.receive()
62+
.doOnError(th -> gatewayHandler.onSessionError(session, th))
5963
.subscribe(
6064
byteBuf ->
61-
Mono.fromCallable(() -> messageCodec.decode(byteBuf))
62-
.doOnNext(message -> metrics.markRequest())
63-
.map(this::checkSid)
64-
.flatMap(msg -> handleCancel(session, msg))
65-
.map(msg -> validateSid(session, (GatewayMessage) msg))
66-
.map(this::checkQualifier)
67-
.map(msg -> gatewayHandler.mapMessage(session, msg))
68-
.subscribe(
69-
request -> {
70-
try {
71-
handleMessage(session, request);
72-
} catch (Exception ex) {
73-
gatewayHandler.onError(session, ex, request, null);
74-
}
75-
},
76-
th -> handleError(session, th)),
77-
th -> gatewayHandler.onError(session, th, null, null));
65+
Mono.deferWithContext(context -> onRequest(session, byteBuf, context))
66+
.subscriberContext(
67+
context -> gatewayHandler.onRequest(session, byteBuf, context))
68+
.subscribe());
7869

7970
return session.onClose(() -> gatewayHandler.onSessionClose(session));
8071
}
8172

82-
private void handleMessage(WebsocketGatewaySession session, GatewayMessage request) {
83-
Long sid = request.streamId();
73+
private Mono<GatewayMessage> onRequest(
74+
WebsocketGatewaySession session, ByteBuf byteBuf, Context context) {
75+
return Mono.fromCallable(() -> messageCodec.decode(byteBuf))
76+
.doOnNext(message -> metrics.markRequest())
77+
.map(this::validateSid)
78+
.flatMap(msg -> onCancel(session, msg))
79+
.map(msg -> validateSid(session, (GatewayMessage) msg))
80+
.map(this::validateQualifier)
81+
.map(msg -> gatewayHandler.mapMessage(session, msg))
82+
.doOnNext(request -> onMessage(session, request, context))
83+
.doOnError(
84+
th -> {
85+
if (!(th instanceof WebsocketContextException)) {
86+
// decode failed at this point
87+
gatewayHandler.onError(session, th, context);
88+
return;
89+
}
90+
91+
WebsocketContextException wex = (WebsocketContextException) th;
92+
wex.releaseRequest(); // release
93+
94+
onError(session, wex.request(), wex.getCause(), context);
95+
});
96+
}
8497

85-
AtomicBoolean receivedError = new AtomicBoolean(false);
98+
private void onMessage(WebsocketGatewaySession session, GatewayMessage request, Context context) {
99+
final Long sid = request.streamId();
100+
final AtomicBoolean receivedError = new AtomicBoolean(false);
86101

87102
final Flux<ServiceMessage> serviceStream =
88103
serviceCall.requestMany(GatewayMessage.toServiceMessage(request));
@@ -94,68 +109,41 @@ private void handleMessage(WebsocketGatewaySession session, GatewayMessage reque
94109
.map(response -> prepareResponse(sid, response, receivedError))
95110
.doOnNext(response -> metrics.markServiceResponse())
96111
.doFinally(signalType -> session.dispose(sid))
97-
.subscribe(
98-
response ->
99-
session
100-
.send(response)
101-
.subscribe(
102-
avoid -> metrics.markResponse(),
103-
th -> gatewayHandler.onError(session, th, request, response)),
104-
th -> handleError(session, request, th),
105-
() -> handleCompletion(session, request, receivedError));
112+
.flatMap(session::send)
113+
.doOnError(th -> onError(session, request, th, context))
114+
.doOnComplete(() -> onComplete(session, request, receivedError, context))
115+
.subscriberContext(context)
116+
.subscribe();
106117

107118
session.register(sid, disposable);
108119
}
109120

110-
private void handleError(WebsocketGatewaySession session, Throwable throwable) {
111-
if (throwable instanceof WebsocketRequestException) {
112-
WebsocketRequestException ex = (WebsocketRequestException) throwable;
113-
ex.releaseRequest(); // release
114-
handleError(session, ex.request(), ex.getCause());
115-
} else {
116-
gatewayHandler.onError(session, throwable, null, null);
117-
}
118-
}
119-
120-
private void handleError(WebsocketGatewaySession session, GatewayMessage req, Throwable th) {
121-
gatewayHandler.onError(session, th, req, null);
121+
private void onError(
122+
WebsocketGatewaySession session, GatewayMessage req, Throwable th, Context context) {
122123

123124
Builder builder = GatewayMessage.from(DefaultErrorMapper.INSTANCE.toMessage(th));
124125
Optional.ofNullable(req.streamId()).ifPresent(builder::streamId);
125126
GatewayMessage response = builder.signal(Signal.ERROR).build();
126127

127-
session
128-
.send(response)
129-
.subscribe(null, ex -> gatewayHandler.onError(session, ex, req, response));
128+
session.send(response).subscriberContext(context).subscribe();
130129
}
131130

132-
private void handleCompletion(
133-
WebsocketGatewaySession session, GatewayMessage req, AtomicBoolean receivedError) {
131+
private void onComplete(
132+
WebsocketGatewaySession session,
133+
GatewayMessage req,
134+
AtomicBoolean receivedError,
135+
Context context) {
136+
134137
if (!receivedError.get()) {
135138
Builder builder = GatewayMessage.builder();
136139
Optional.ofNullable(req.streamId()).ifPresent(builder::streamId);
137140
GatewayMessage response = builder.signal(Signal.COMPLETE).build();
138-
session.send(response).subscribe(null, ex -> gatewayHandler.onError(session, ex, req, null));
139-
}
140-
}
141141

142-
private GatewayMessage checkQualifier(GatewayMessage msg) {
143-
if (msg.qualifier() == null) {
144-
throw WebsocketRequestException.newBadRequest("qualifier is missing", msg);
142+
session.send(response).subscriberContext(context).subscribe();
145143
}
146-
return msg;
147144
}
148145

149-
private GatewayMessage validateSid(WebsocketGatewaySession session, GatewayMessage msg) {
150-
if (session.containsSid(msg.streamId())) {
151-
throw WebsocketRequestException.newBadRequest(
152-
"sid=" + msg.streamId() + " is already registered", msg);
153-
} else {
154-
return msg;
155-
}
156-
}
157-
158-
private Mono<?> handleCancel(WebsocketGatewaySession session, GatewayMessage msg) {
146+
private Mono<?> onCancel(WebsocketGatewaySession session, GatewayMessage msg) {
159147
if (!msg.hasSignal(Signal.CANCEL)) {
160148
return Mono.just(msg);
161149
}
@@ -171,9 +159,25 @@ private Mono<?> handleCancel(WebsocketGatewaySession session, GatewayMessage msg
171159
return session.send(cancelAck); // no need to subscribe here since flatMap will do
172160
}
173161

174-
private GatewayMessage checkSid(GatewayMessage msg) {
162+
private GatewayMessage validateQualifier(GatewayMessage msg) {
163+
if (msg.qualifier() == null) {
164+
throw WebsocketContextException.badRequest("qualifier is missing", msg);
165+
}
166+
return msg;
167+
}
168+
169+
private GatewayMessage validateSid(WebsocketGatewaySession session, GatewayMessage msg) {
170+
if (session.containsSid(msg.streamId())) {
171+
throw WebsocketContextException.badRequest(
172+
"sid=" + msg.streamId() + " is already registered", msg);
173+
} else {
174+
return msg;
175+
}
176+
}
177+
178+
private GatewayMessage validateSid(GatewayMessage msg) {
175179
if (msg.streamId() == null) {
176-
throw WebsocketRequestException.newBadRequest("sid is missing", msg);
180+
throw WebsocketContextException.badRequest("sid is missing", msg);
177181
} else {
178182
return msg;
179183
}

0 commit comments

Comments
 (0)