Skip to content

Commit a1a8781

Browse files
committed
PayloadUtils improvements and tests
1 parent 22e87ac commit a1a8781

File tree

3 files changed

+207
-34
lines changed

3 files changed

+207
-34
lines changed

spring-messaging/src/main/java/org/springframework/messaging/rsocket/MetadataEncoder.java

Lines changed: 3 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -24,14 +24,12 @@
2424

2525
import io.netty.buffer.ByteBufAllocator;
2626
import io.netty.buffer.CompositeByteBuf;
27-
import io.netty.buffer.Unpooled;
2827
import io.rsocket.metadata.CompositeMetadataFlyweight;
2928

3029
import org.springframework.core.ResolvableType;
3130
import org.springframework.core.codec.Encoder;
3231
import org.springframework.core.io.buffer.DataBuffer;
3332
import org.springframework.core.io.buffer.DataBufferFactory;
34-
import org.springframework.core.io.buffer.NettyDataBuffer;
3533
import org.springframework.core.io.buffer.NettyDataBufferFactory;
3634
import org.springframework.lang.Nullable;
3735
import org.springframework.util.Assert;
@@ -165,18 +163,14 @@ public DataBuffer encode() {
165163
try {
166164
mergedMetadata.forEach((value, mimeType) -> {
167165
DataBuffer buffer = encodeEntry(value, mimeType);
168-
CompositeMetadataFlyweight.encodeAndAddMetadata(composite, this.allocator,
169-
mimeType.toString(),
170-
buffer instanceof NettyDataBuffer ?
171-
((NettyDataBuffer) buffer).getNativeBuffer() :
172-
Unpooled.wrappedBuffer(buffer.asByteBuffer()));
166+
CompositeMetadataFlyweight.encodeAndAddMetadata(
167+
composite, this.allocator, mimeType.toString(), PayloadUtils.asByteBuf(buffer));
173168
});
174169
if (bufferFactory() instanceof NettyDataBufferFactory) {
175170
return ((NettyDataBufferFactory) bufferFactory()).wrap(composite);
176171
}
177172
else {
178-
DataBuffer buffer = bufferFactory().allocateBuffer();
179-
buffer.write(composite.nioBuffer());
173+
DataBuffer buffer = bufferFactory().wrap(composite.nioBuffer());
180174
composite.release();
181175
return buffer;
182176
}

spring-messaging/src/main/java/org/springframework/messaging/rsocket/PayloadUtils.java

Lines changed: 29 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,10 @@
1616

1717
package org.springframework.messaging.rsocket;
1818

19+
import java.nio.ByteBuffer;
20+
1921
import io.netty.buffer.ByteBuf;
22+
import io.netty.buffer.Unpooled;
2023
import io.rsocket.Payload;
2124
import io.rsocket.util.ByteBufPayload;
2225
import io.rsocket.util.DefaultPayload;
@@ -64,41 +67,42 @@ public static DataBuffer retainDataAndReleasePayload(Payload payload, DataBuffer
6467

6568
/**
6669
* Create a Payload from the given metadata and data.
70+
* <p>If at least one is {@link NettyDataBuffer} then {@link ByteBufPayload}
71+
* is created with either obtaining the underlying native {@link ByteBuf}
72+
* or using {@link Unpooled#wrappedBuffer(ByteBuffer...)} if necessary.
73+
* Otherwise, if both are {@link DefaultDataBuffer}, then
74+
* {@link DefaultPayload} is created.
6775
* @param metadata the metadata part for the payload
6876
* @param data the data part for the payload
69-
* @return the created Payload
77+
* @return the created payload
7078
*/
7179
public static Payload createPayload(DataBuffer metadata, DataBuffer data) {
72-
if (metadata instanceof NettyDataBuffer && data instanceof NettyDataBuffer) {
73-
return ByteBufPayload.create(
74-
((NettyDataBuffer) data).getNativeBuffer(),
75-
((NettyDataBuffer) metadata).getNativeBuffer());
76-
}
77-
else if (metadata instanceof DefaultDataBuffer && data instanceof DefaultDataBuffer) {
78-
return DefaultPayload.create(
79-
((DefaultDataBuffer) data).getNativeBuffer(),
80-
((DefaultDataBuffer) metadata).getNativeBuffer());
81-
}
82-
else {
83-
return DefaultPayload.create(data.asByteBuffer(), metadata.asByteBuffer());
84-
}
80+
return data instanceof NettyDataBuffer || metadata instanceof NettyDataBuffer ?
81+
ByteBufPayload.create(asByteBuf(data), asByteBuf(metadata)) :
82+
DefaultPayload.create(asByteBuffer(data), asByteBuffer(metadata));
8583
}
8684

8785
/**
88-
* Create a Payload from the given data.
86+
* Create a Payload with data only. The created payload is
87+
* {@link ByteBufPayload} if the input is {@link NettyDataBuffer} or
88+
* otherwise it is {@link DefaultPayload}.
8989
* @param data the data part for the payload
90-
* @return the created Payload
90+
* @return created payload
9191
*/
9292
public static Payload createPayload(DataBuffer data) {
93-
if (data instanceof NettyDataBuffer) {
94-
return ByteBufPayload.create(((NettyDataBuffer) data).getNativeBuffer());
95-
}
96-
else if (data instanceof DefaultDataBuffer) {
97-
return DefaultPayload.create(((DefaultDataBuffer) data).getNativeBuffer());
98-
}
99-
else {
100-
return DefaultPayload.create(data.asByteBuffer());
101-
}
93+
return data instanceof NettyDataBuffer ?
94+
ByteBufPayload.create(asByteBuf(data)) : DefaultPayload.create(asByteBuffer(data));
95+
}
96+
97+
98+
static ByteBuf asByteBuf(DataBuffer buffer) {
99+
return buffer instanceof NettyDataBuffer ?
100+
((NettyDataBuffer) buffer).getNativeBuffer() : Unpooled.wrappedBuffer(buffer.asByteBuffer());
101+
}
102+
103+
private static ByteBuffer asByteBuffer(DataBuffer buffer) {
104+
return buffer instanceof DefaultDataBuffer ?
105+
((DefaultDataBuffer) buffer).getNativeBuffer() : buffer.asByteBuffer();
102106
}
103107

104108
}
Lines changed: 175 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,175 @@
1+
/*
2+
* Copyright 2002-2019 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+
* https://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+
package org.springframework.messaging.rsocket;
17+
18+
import java.nio.charset.StandardCharsets;
19+
import java.time.Duration;
20+
21+
import io.netty.buffer.PooledByteBufAllocator;
22+
import io.rsocket.Payload;
23+
import io.rsocket.util.ByteBufPayload;
24+
import io.rsocket.util.DefaultPayload;
25+
import org.junit.After;
26+
import org.junit.Test;
27+
28+
import org.springframework.core.io.buffer.DataBuffer;
29+
import org.springframework.core.io.buffer.DataBufferUtils;
30+
import org.springframework.core.io.buffer.DefaultDataBuffer;
31+
import org.springframework.core.io.buffer.DefaultDataBufferFactory;
32+
import org.springframework.core.io.buffer.NettyDataBuffer;
33+
import org.springframework.core.io.buffer.support.DataBufferTestUtils;
34+
35+
import static org.assertj.core.api.Assertions.assertThat;
36+
37+
/**
38+
* Unit tests for {@link PayloadUtils}.
39+
* @author Rossen Stoyanchev
40+
* @since 5.2
41+
*/
42+
public class PayloadUtilsTests {
43+
44+
private LeakAwareNettyDataBufferFactory nettyBufferFactory =
45+
new LeakAwareNettyDataBufferFactory(PooledByteBufAllocator.DEFAULT);
46+
47+
private DefaultDataBufferFactory defaultBufferFactory = new DefaultDataBufferFactory();
48+
49+
50+
@After
51+
public void tearDown() throws Exception {
52+
this.nettyBufferFactory.checkForLeaks(Duration.ofSeconds(5));
53+
}
54+
55+
56+
@Test
57+
public void retainAndReleaseWithNettyFactory() {
58+
Payload payload = ByteBufPayload.create("sample data");
59+
DataBuffer buffer = PayloadUtils.retainDataAndReleasePayload(payload, this.nettyBufferFactory);
60+
try {
61+
assertThat(buffer).isInstanceOf(NettyDataBuffer.class);
62+
assertThat(((NettyDataBuffer) buffer).getNativeBuffer().refCnt()).isEqualTo(1);
63+
assertThat(payload.refCnt()).isEqualTo(0);
64+
}
65+
finally {
66+
DataBufferUtils.release(buffer);
67+
}
68+
}
69+
70+
@Test
71+
public void retainAndReleaseWithDefaultFactory() {
72+
Payload payload = ByteBufPayload.create("sample data");
73+
DataBuffer buffer = PayloadUtils.retainDataAndReleasePayload(payload, this.defaultBufferFactory);
74+
75+
assertThat(buffer).isInstanceOf(DefaultDataBuffer.class);
76+
assertThat(payload.refCnt()).isEqualTo(0);
77+
}
78+
79+
@Test
80+
public void createWithNettyBuffers() {
81+
NettyDataBuffer data = createNettyDataBuffer("sample data");
82+
NettyDataBuffer metadata = createNettyDataBuffer("sample metadata");
83+
84+
Payload payload = PayloadUtils.createPayload(metadata, data);
85+
try {
86+
assertThat(payload).isInstanceOf(ByteBufPayload.class);
87+
assertThat(payload.data()).isSameAs(data.getNativeBuffer());
88+
assertThat(payload.metadata()).isSameAs(metadata.getNativeBuffer());
89+
}
90+
finally {
91+
payload.release();
92+
}
93+
}
94+
95+
@Test
96+
public void createWithDefaultBuffers() {
97+
DataBuffer data = createDefaultDataBuffer("sample data");
98+
DataBuffer metadata = createDefaultDataBuffer("sample metadata");
99+
Payload payload = PayloadUtils.createPayload(metadata, data);
100+
101+
assertThat(payload).isInstanceOf(DefaultPayload.class);
102+
assertThat(payload.getDataUtf8()).isEqualTo(dataBufferToString(data));
103+
assertThat(payload.getMetadataUtf8()).isEqualTo(dataBufferToString(metadata));
104+
}
105+
106+
@Test
107+
public void createWithNettyAndDefaultBuffers() {
108+
NettyDataBuffer data = createNettyDataBuffer("sample data");
109+
DefaultDataBuffer metadata = createDefaultDataBuffer("sample metadata");
110+
Payload payload = PayloadUtils.createPayload(metadata, data);
111+
try {
112+
assertThat(payload).isInstanceOf(ByteBufPayload.class);
113+
assertThat(payload.data()).isSameAs(data.getNativeBuffer());
114+
assertThat(payload.getMetadataUtf8()).isEqualTo(dataBufferToString(metadata));
115+
}
116+
finally {
117+
payload.release();
118+
}
119+
}
120+
121+
@Test
122+
public void createWithDefaultAndNettyBuffers() {
123+
DefaultDataBuffer data = createDefaultDataBuffer("sample data");
124+
NettyDataBuffer metadata = createNettyDataBuffer("sample metadata");
125+
Payload payload = PayloadUtils.createPayload(metadata, data);
126+
try {
127+
assertThat(payload).isInstanceOf(ByteBufPayload.class);
128+
assertThat(payload.getDataUtf8()).isEqualTo(dataBufferToString(data));
129+
assertThat(payload.metadata()).isSameAs(metadata.getNativeBuffer());
130+
}
131+
finally {
132+
payload.release();
133+
}
134+
}
135+
136+
@Test
137+
public void createWithNettyBuffer() {
138+
NettyDataBuffer data = createNettyDataBuffer("sample data");
139+
Payload payload = PayloadUtils.createPayload(data);
140+
try {
141+
assertThat(payload).isInstanceOf(ByteBufPayload.class);
142+
assertThat(payload.data()).isSameAs(data.getNativeBuffer());
143+
}
144+
finally {
145+
payload.release();
146+
}
147+
}
148+
149+
@Test
150+
public void createWithDefaultBuffer() {
151+
DataBuffer data = createDefaultDataBuffer("sample data");
152+
Payload payload = PayloadUtils.createPayload(data);
153+
154+
assertThat(payload).isInstanceOf(DefaultPayload.class);
155+
assertThat(payload.getDataUtf8()).isEqualTo(dataBufferToString(data));
156+
}
157+
158+
159+
private NettyDataBuffer createNettyDataBuffer(String content) {
160+
NettyDataBuffer buffer = this.nettyBufferFactory.allocateBuffer();
161+
buffer.write(content, StandardCharsets.UTF_8);
162+
return buffer;
163+
}
164+
165+
private DefaultDataBuffer createDefaultDataBuffer(String content) {
166+
DefaultDataBuffer buffer = this.defaultBufferFactory.allocateBuffer();
167+
buffer.write(content, StandardCharsets.UTF_8);
168+
return buffer;
169+
}
170+
171+
private String dataBufferToString(DataBuffer metadata) {
172+
return DataBufferTestUtils.dumpString(metadata, StandardCharsets.UTF_8);
173+
}
174+
175+
}

0 commit comments

Comments
 (0)