Skip to content

Commit 6c82ecc

Browse files
connieyconniey
authored
Adds support for Continuation WebSocketFrames (Azure#44130)
* Cast Consumer to WebPubSubMessage. * Add log message for unknown message type. * Update Netty implementation to follow interface. * Update Consumer for WebPubSubMessage. * Update Consumer for WebPubSubMessage in test. * Decode concrete WebPubSubMessage. * Support continuation web socket frames. * Update echo sample. * Updating pom.xml for tests. * Add tests * formatting changes from spotless * Disposing of frame when finished. * Add changelog * Revert usage of new API. * Update to use COmpositeByteBuffer * Checkstyle changes * Update write index when adding buffer. * Update CHANGELOG * Release on exception. * Add guard for releasing bytebuffer and check of there are no buffers to publish. * Fix CHANGELOG.md * Remove duplicate code. * Update dep versions
1 parent 141e42e commit 6c82ecc

File tree

11 files changed

+375
-73
lines changed

11 files changed

+375
-73
lines changed

‎sdk/webpubsub/azure-messaging-webpubsub-client/CHANGELOG.md‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,8 @@
1313

1414
### Bugs Fixed
1515

16+
- Fixes issue where text spanning more than a single WebSocketFrame are ignored. [#44130](https://github.com/Azure/azure-sdk-for-java/pull/44130)
17+
1618
### Other Changes
1719

1820

‎sdk/webpubsub/azure-messaging-webpubsub-client/pom.xml‎

Lines changed: 16 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,10 @@
3333
</scm>
3434

3535
<properties>
36+
<javaModulesSurefireArgLine>
37+
--add-exports com.azure.core/com.azure.core.implementation=ALL-UNNAMED
38+
</javaModulesSurefireArgLine>
39+
3640
<!-- sdk uses live tests -->
3741
<jacoco.min.linecoverage>0.2</jacoco.min.linecoverage>
3842
<jacoco.min.branchcoverage>0.1</jacoco.min.branchcoverage>
@@ -63,6 +67,12 @@
6367
<version>1.27.0-beta.6</version> <!-- {x-version-update;com.azure:azure-core-test;dependency} -->
6468
<scope>test</scope>
6569
</dependency>
70+
<dependency>
71+
<groupId>com.azure</groupId>
72+
<artifactId>azure-core-serializer-json-jackson</artifactId>
73+
<version>1.5.6</version> <!-- {x-version-update;com.azure:azure-core-serializer-json-jackson;dependency} -->
74+
<scope>test</scope>
75+
</dependency>
6676
<dependency>
6777
<groupId>com.azure</groupId>
6878
<artifactId>azure-identity</artifactId>
@@ -75,6 +85,12 @@
7585
<version>1.4.0-beta.1</version> <!-- {x-version-update;com.azure:azure-messaging-webpubsub;current} -->
7686
<scope>test</scope>
7787
</dependency>
88+
<dependency>
89+
<groupId>org.mockito</groupId>
90+
<artifactId>mockito-core</artifactId>
91+
<version>4.11.0</version> <!-- {x-version-update;org.mockito:mockito-core;external_dependency} -->
92+
<scope>test</scope>
93+
</dependency>
7894
</dependencies>
7995

8096
<build>

‎sdk/webpubsub/azure-messaging-webpubsub-client/src/main/java/com/azure/messaging/webpubsub/client/WebPubSubAsyncClient.java‎

Lines changed: 10 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -529,16 +529,6 @@ private Flux<AckMessage> receiveAckMessages() {
529529

530530
private Mono<Void> sendMessage(WebPubSubMessage message) {
531531
return checkStateBeforeSend().then(Mono.create(sink -> {
532-
// if (logger.canLogAtLevel(LogLevel.VERBOSE)) {
533-
// try {
534-
// String json = JacksonAdapter.createDefaultSerializerAdapter()
535-
// .serialize(message, SerializerEncoding.JSON);
536-
// logger.atVerbose().addKeyValue("message", json).log("Send message");
537-
// } catch (IOException e) {
538-
// sink.error(new UncheckedIOException("Failed to serialize message for VERBOSE logging", e));
539-
// }
540-
// }
541-
542532
webSocketSession.sendObjectAsync(message, sendResult -> {
543533
if (sendResult.isOK()) {
544534
sink.success();
@@ -758,18 +748,7 @@ private void handleSessionClose(CloseReason closeReason) {
758748
}
759749
}
760750

761-
private void handleMessage(Object webPubSubMessage) {
762-
// if (logger.canLogAtLevel(LogLevel.VERBOSE)) {
763-
// try {
764-
// String json = JacksonAdapter.createDefaultSerializerAdapter()
765-
// .serialize(webPubSubMessage, SerializerEncoding.JSON);
766-
// logger.atVerbose().addKeyValue("message", json).log("Received message");
767-
// } catch (IOException e) {
768-
// throw logger.logExceptionAsError(
769-
// new UncheckedIOException("Failed to serialize received message for VERBOSE logging", e));
770-
// }
771-
// }
772-
751+
private void handleMessage(WebPubSubMessage webPubSubMessage) {
773752
if (webPubSubMessage instanceof GroupDataMessage) {
774753
final GroupDataMessage groupDataMessage = (GroupDataMessage) webPubSubMessage;
775754

@@ -783,6 +762,7 @@ private void handleMessage(Object webPubSubMessage) {
783762
groupDataMessage.getDataType(), groupDataMessage.getFromUserId(),
784763
groupDataMessage.getSequenceId()));
785764
}
765+
786766
} else if (webPubSubMessage instanceof ServerDataMessage) {
787767
final ServerDataMessage serverDataMessage = (ServerDataMessage) webPubSubMessage;
788768

@@ -815,6 +795,14 @@ private void handleMessage(Object webPubSubMessage) {
815795
final DisconnectedMessage disconnectedMessage = (DisconnectedMessage) webPubSubMessage;
816796
// send DisconnectedEvent, but connection close will be handled in handleSessionClose
817797
handleConnectionClose(new DisconnectedEvent(this.getConnectionId(), disconnectedMessage.getReason()));
798+
} else {
799+
final ClientLogger logger = loggerReference.get();
800+
if (logger != null) {
801+
logger.atWarning()
802+
.addKeyValue("type", webPubSubMessage.getClass())
803+
.addKeyValue("message", webPubSubMessage)
804+
.log("Unknown message type. Skipping decode.");
805+
}
818806
}
819807
}
820808

‎sdk/webpubsub/azure-messaging-webpubsub-client/src/main/java/com/azure/messaging/webpubsub/client/implementation/MessageDecoder.java‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -11,7 +11,7 @@
1111
import java.io.UncheckedIOException;
1212

1313
public final class MessageDecoder {
14-
public Object decode(String s) {
14+
public WebPubSubMessage decode(String s) {
1515
try (JsonReader jsonReader = JsonProviders.createReader(s)) {
1616
return WebPubSubMessage.fromJson(jsonReader);
1717
} catch (IOException e) {

‎sdk/webpubsub/azure-messaging-webpubsub-client/src/main/java/com/azure/messaging/webpubsub/client/implementation/websocket/WebSocketClient.java‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package com.azure.messaging.webpubsub.client.implementation.websocket;
55

66
import com.azure.core.util.logging.ClientLogger;
7+
import com.azure.messaging.webpubsub.client.implementation.models.WebPubSubMessage;
78

89
import java.util.concurrent.atomic.AtomicReference;
910
import java.util.function.Consumer;
@@ -12,6 +13,6 @@
1213
public interface WebSocketClient {
1314

1415
WebSocketSession connectToServer(ClientEndpointConfiguration cec, String path,
15-
AtomicReference<ClientLogger> loggerReference, Consumer<Object> messageHandler,
16+
AtomicReference<ClientLogger> loggerReference, Consumer<WebPubSubMessage> messageHandler,
1617
Consumer<WebSocketSession> openHandler, Consumer<CloseReason> closeHandler);
1718
}

‎sdk/webpubsub/azure-messaging-webpubsub-client/src/main/java/com/azure/messaging/webpubsub/client/implementation/websocket/WebSocketClientHandler.java‎

Lines changed: 104 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -3,38 +3,45 @@
33

44
package com.azure.messaging.webpubsub.client.implementation.websocket;
55

6+
import com.azure.core.util.BinaryData;
67
import com.azure.core.util.logging.ClientLogger;
78
import com.azure.messaging.webpubsub.client.implementation.MessageDecoder;
9+
import com.azure.messaging.webpubsub.client.implementation.models.WebPubSubMessage;
10+
import io.netty.buffer.ByteBufAllocator;
11+
import io.netty.buffer.CompositeByteBuf;
812
import io.netty.channel.Channel;
913
import io.netty.channel.ChannelFuture;
1014
import io.netty.channel.ChannelHandlerContext;
1115
import io.netty.channel.ChannelPromise;
1216
import io.netty.channel.SimpleChannelInboundHandler;
1317
import io.netty.handler.codec.http.FullHttpResponse;
1418
import io.netty.handler.codec.http.websocketx.CloseWebSocketFrame;
19+
import io.netty.handler.codec.http.websocketx.ContinuationWebSocketFrame;
1520
import io.netty.handler.codec.http.websocketx.PingWebSocketFrame;
1621
import io.netty.handler.codec.http.websocketx.PongWebSocketFrame;
1722
import io.netty.handler.codec.http.websocketx.TextWebSocketFrame;
1823
import io.netty.handler.codec.http.websocketx.WebSocketClientHandshaker;
1924
import io.netty.handler.codec.http.websocketx.WebSocketFrame;
2025
import io.netty.handler.codec.http.websocketx.WebSocketHandshakeException;
21-
import io.netty.util.CharsetUtil;
2226

27+
import java.nio.ByteBuffer;
28+
import java.util.Arrays;
2329
import java.util.concurrent.CompletableFuture;
2430
import java.util.concurrent.atomic.AtomicReference;
2531
import java.util.function.Consumer;
2632

2733
final class WebSocketClientHandler extends SimpleChannelInboundHandler<Object> {
2834

2935
private final WebSocketClientHandshaker handshaker;
30-
private ChannelPromise handshakeFuture;
31-
3236
private final AtomicReference<ClientLogger> loggerReference;
3337
private final MessageDecoder messageDecoder;
34-
private final Consumer<Object> messageHandler;
38+
private final Consumer<WebPubSubMessage> messageHandler;
39+
40+
private ChannelPromise handshakeFuture;
41+
private CompositeByteBuf compositeByteBuf;
3542

3643
WebSocketClientHandler(WebSocketClientHandshaker handshaker, AtomicReference<ClientLogger> loggerReference,
37-
MessageDecoder messageDecoder, Consumer<Object> messageHandler) {
44+
MessageDecoder messageDecoder, Consumer<WebPubSubMessage> messageHandler) {
3845
this.handshaker = handshaker;
3946
this.loggerReference = loggerReference;
4047
this.messageDecoder = messageDecoder;
@@ -46,80 +53,131 @@ ChannelFuture handshakeFuture() {
4653
}
4754

4855
@Override
49-
public void handlerAdded(ChannelHandlerContext ctx) {
50-
handshakeFuture = ctx.newPromise();
56+
public void handlerAdded(ChannelHandlerContext context) {
57+
handshakeFuture = context.newPromise();
58+
compositeByteBuf = context.alloc().compositeBuffer();
59+
}
60+
61+
@Override
62+
public void handlerRemoved(ChannelHandlerContext ctx) {
63+
publishBuffer();
5164
}
5265

5366
@Override
54-
public void channelActive(ChannelHandlerContext ctx) {
55-
handshaker.handshake(ctx.channel());
67+
public void channelActive(ChannelHandlerContext context) {
68+
handshaker.handshake(context.channel());
5669
}
5770

5871
@Override
59-
protected void channelRead0(ChannelHandlerContext ctx, Object msg) {
60-
Channel ch = ctx.channel();
72+
protected void channelRead0(ChannelHandlerContext context, Object message) {
6173
if (handshakeFuture != null && !handshaker.isHandshakeComplete()) {
74+
Channel channel = context.channel();
75+
6276
try {
63-
handshaker.finishHandshake(ch, (FullHttpResponse) msg);
77+
handshaker.finishHandshake(channel, (FullHttpResponse) message);
6478
handshakeFuture.setSuccess();
6579
} catch (WebSocketHandshakeException e) {
6680
handshakeFuture.setFailure(e);
6781
}
6882
return;
6983
}
7084

71-
if (msg instanceof FullHttpResponse) {
72-
FullHttpResponse response = (FullHttpResponse) msg;
73-
throw loggerReference.get()
74-
.logExceptionAsError(new IllegalStateException("Unexpected FullHttpResponse (getStatus="
75-
+ response.status() + ", content=" + response.content().toString(CharsetUtil.UTF_8) + ')'));
85+
if (!(message instanceof WebSocketFrame) || !processMessage(context, (WebSocketFrame) message)) {
86+
loggerReference.get()
87+
.atWarning()
88+
.addKeyValue("messageType", message.getClass())
89+
.log("Unknown message type. Skipping.");
90+
91+
context.fireChannelRead(message);
92+
}
93+
}
94+
95+
@Override
96+
public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
97+
if (handshakeFuture != null && !handshakeFuture.isDone()) {
98+
handshakeFuture.setFailure(cause);
7699
}
100+
ctx.close();
101+
release(compositeByteBuf);
102+
}
77103

78-
WebSocketFrame frame = (WebSocketFrame) msg;
79-
if (frame instanceof TextWebSocketFrame) {
80-
// Text
81-
TextWebSocketFrame textFrame = (TextWebSocketFrame) frame;
82-
loggerReference.get().atVerbose().addKeyValue("text", textFrame.text()).log("Received TextWebSocketFrame");
83-
Object wpsMessage = messageDecoder.decode(textFrame.text());
84-
messageHandler.accept(wpsMessage);
85-
} else if (frame instanceof PingWebSocketFrame) {
104+
/**
105+
* Attempts to process the web socket frame.
106+
*
107+
* @param context Channel for message.
108+
* @param webSocketFrame Frame to process.
109+
* @return true if the frame was processed, false otherwise.
110+
*/
111+
private boolean processMessage(ChannelHandlerContext context, WebSocketFrame webSocketFrame) {
112+
Channel channel = context.channel();
113+
114+
if (webSocketFrame instanceof PingWebSocketFrame) {
86115
// Ping, reply Pong
87116
loggerReference.get().atVerbose().log("Received PingWebSocketFrame");
88117
loggerReference.get().atVerbose().log("Send PongWebSocketFrame");
89-
ch.writeAndFlush(new PongWebSocketFrame());
90-
} else if (frame instanceof PongWebSocketFrame) {
118+
channel.writeAndFlush(new PongWebSocketFrame());
119+
return true;
120+
} else if (webSocketFrame instanceof PongWebSocketFrame) {
91121
// Pong
92122
loggerReference.get().atVerbose().log("Received PongWebSocketFrame");
93-
} else if (frame instanceof CloseWebSocketFrame) {
123+
return true;
124+
} else if (webSocketFrame instanceof CloseWebSocketFrame) {
94125
// Close
95-
CloseWebSocketFrame closeFrame = (CloseWebSocketFrame) frame;
126+
final CloseWebSocketFrame closeFrame = (CloseWebSocketFrame) webSocketFrame;
127+
96128
loggerReference.get()
97129
.atVerbose()
98130
.addKeyValue("statusCode", closeFrame.statusCode())
99131
.addKeyValue("reasonText", closeFrame.reasonText())
100132
.log("Received CloseWebSocketFrame");
101-
102-
this.serverCloseWebSocketFrame = closeFrame.retain(); // retain for SessionNettyImpl
133+
serverCloseWebSocketFrame = closeFrame.retain(); // retain for SessionNettyImpl
103134

104135
if (closeCallbackFuture == null) {
105136
// close initiated from server, reply CloseWebSocketFrame, then close connection
106137
loggerReference.get().atVerbose().log("Send CloseWebSocketFrame");
107138
closeFrame.retain(); // retain before write it back
108-
ch.writeAndFlush(closeFrame).addListener(future -> ch.close());
139+
channel.writeAndFlush(closeFrame).addListener(future -> channel.close());
109140
} else {
110141
// close initiated from client, client already sent CloseWebSocketFrame
111-
ch.close();
142+
channel.close();
143+
}
144+
145+
return true;
146+
} else if (webSocketFrame instanceof TextWebSocketFrame
147+
|| webSocketFrame instanceof ContinuationWebSocketFrame) {
148+
if (compositeByteBuf == null) {
149+
compositeByteBuf = ByteBufAllocator.DEFAULT.compositeBuffer();
150+
}
151+
152+
compositeByteBuf.addComponent(true, webSocketFrame.content().retain());
153+
154+
if (!webSocketFrame.isFinalFragment()) {
155+
return true;
112156
}
157+
158+
publishBuffer();
159+
return true;
160+
} else {
161+
return false;
113162
}
114163
}
115164

116-
@Override
117-
public void exceptionCaught(ChannelHandlerContext ctx, Throwable cause) {
118-
cause.printStackTrace();
119-
if (handshakeFuture != null && !handshakeFuture.isDone()) {
120-
handshakeFuture.setFailure(cause);
165+
private void publishBuffer() {
166+
final ByteBuffer[] nioBuffers = compositeByteBuf.nioBuffers();
167+
168+
if (nioBuffers.length == 0) {
169+
return;
170+
}
171+
172+
try {
173+
final BinaryData data = BinaryData.fromListByteBuffer(Arrays.asList(nioBuffers));
174+
final String collected = data.toString();
175+
final WebPubSubMessage deserialized = messageDecoder.decode(collected);
176+
177+
messageHandler.accept(deserialized);
178+
} finally {
179+
release(compositeByteBuf);
121180
}
122-
ctx.close();
123181
}
124182

125183
// as side effect, if it is not null, the close (aka CloseWebSocketFrame) is initiated by client
@@ -138,4 +196,11 @@ public CompletableFuture<Void> getClientCloseCallbackFuture() {
138196
CloseWebSocketFrame getServerCloseWebSocketFrame() {
139197
return this.serverCloseWebSocketFrame;
140198
}
199+
200+
private static void release(CompositeByteBuf buffer) {
201+
if (buffer.refCnt() > 0) {
202+
buffer.release();
203+
buffer.clear();
204+
}
205+
}
141206
}

‎sdk/webpubsub/azure-messaging-webpubsub-client/src/main/java/com/azure/messaging/webpubsub/client/implementation/websocket/WebSocketClientNettyImpl.java‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
package com.azure.messaging.webpubsub.client.implementation.websocket;
55

66
import com.azure.core.util.logging.ClientLogger;
7+
import com.azure.messaging.webpubsub.client.implementation.models.WebPubSubMessage;
78
import com.azure.messaging.webpubsub.client.models.ConnectFailedException;
89

910
import java.util.concurrent.atomic.AtomicReference;
@@ -12,12 +13,14 @@
1213
public final class WebSocketClientNettyImpl implements WebSocketClient {
1314
@Override
1415
public WebSocketSession connectToServer(ClientEndpointConfiguration cec, String path,
15-
AtomicReference<ClientLogger> loggerReference, Consumer<Object> messageHandler,
16+
AtomicReference<ClientLogger> loggerReference, Consumer<WebPubSubMessage> messageHandler,
1617
Consumer<WebSocketSession> openHandler, Consumer<CloseReason> closeHandler) {
1718
try {
1819
WebSocketSessionNettyImpl session
1920
= new WebSocketSessionNettyImpl(cec, path, loggerReference, messageHandler, openHandler, closeHandler);
21+
2022
session.connect();
23+
2124
return session;
2225
} catch (Exception e) {
2326
throw loggerReference.get().logExceptionAsError(new ConnectFailedException("Failed to connect", e));

0 commit comments

Comments
 (0)