Skip to content

Commit e83793b

Browse files
bclozelbclozel
committed
Batch SSE events writes when possible
Prior to this commit, the `SseEventBuilder` would be used to create SSE events and write them to the connection using the `ResponseBodyEmitter`. This would send each data item one by one, effectively writing and flushing to the network for each. Since multiple data lines are prepared by the `SseEventBuilder`, a typical write of an SSE event performs multiple flushes operations. This commit adds a method on `ResponseBodyEmitter` to perform batch writes (given a `Set<DataWithMediaType>`) and only flush once all elements of the set have been written. This also applies in case of early writes, where now all buffered elements are written then flushed altogether. Fixes gh-30912
1 parent 18966d0 commit e83793b

File tree

6 files changed

+108
-33
lines changed

6 files changed

+108
-33
lines changed

‎spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitter.java

Lines changed: 48 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -128,9 +128,7 @@ synchronized void initialize(Handler handler) throws IOException {
128128
this.handler = handler;
129129

130130
try {
131-
for (DataWithMediaType sendAttempt : this.earlySendAttempts) {
132-
sendInternal(sendAttempt.getData(), sendAttempt.getMediaType());
133-
}
131+
sendInternal(this.earlySendAttempts);
134132
}
135133
finally {
136134
this.earlySendAttempts.clear();
@@ -194,11 +192,7 @@ public void send(Object object) throws IOException {
194192
*/
195193
public synchronized void send(Object object, @Nullable MediaType mediaType) throws IOException {
196194
Assert.state(!this.complete, () -> "ResponseBodyEmitter has already completed" +
197-
(this.failure != null ? " with error: " + this.failure : ""));
198-
sendInternal(object, mediaType);
199-
}
200-
201-
private void sendInternal(Object object, @Nullable MediaType mediaType) throws IOException {
195+
(this.failure != null ? " with error: " + this.failure : ""));
202196
if (this.handler != null) {
203197
try {
204198
this.handler.send(object, mediaType);
@@ -217,6 +211,43 @@ private void sendInternal(Object object, @Nullable MediaType mediaType) throws I
217211
}
218212
}
219213

214+
/**
215+
* Write a set of data and MediaType pairs in a batch.
216+
* <p>Compared to {@link #send(Object, MediaType)}, this batches the write operations
217+
* and flushes to the network at the end.
218+
* @param items the object and media type pairs to write
219+
* @throws IOException raised when an I/O error occurs
220+
* @throws java.lang.IllegalStateException wraps any other errors
221+
* @since 6.0.12
222+
*/
223+
public synchronized void send(Set<DataWithMediaType> items) throws IOException {
224+
Assert.state(!this.complete, () -> "ResponseBodyEmitter has already completed" +
225+
(this.failure != null ? " with error: " + this.failure : ""));
226+
sendInternal(items);
227+
}
228+
229+
private void sendInternal(Set<DataWithMediaType> items) throws IOException {
230+
if (items.isEmpty()) {
231+
return;
232+
}
233+
if (this.handler != null) {
234+
try {
235+
this.handler.send(items);
236+
}
237+
catch (IOException ex) {
238+
this.sendFailed = true;
239+
throw ex;
240+
}
241+
catch (Throwable ex) {
242+
this.sendFailed = true;
243+
throw new IllegalStateException("Failed to send " + items, ex);
244+
}
245+
}
246+
else {
247+
this.earlySendAttempts.addAll(items);
248+
}
249+
}
250+
220251
/**
221252
* Complete request processing by performing a dispatch into the servlet
222253
* container, where Spring MVC is invoked once more, and completes the
@@ -302,8 +333,17 @@ public String toString() {
302333
*/
303334
interface Handler {
304335

336+
/**
337+
* Immediately write and flush the given data to the network.
338+
*/
305339
void send(Object data, @Nullable MediaType mediaType) throws IOException;
306340

341+
/**
342+
* Immediately write all data items then flush to the network.
343+
* @since 6.0.12
344+
*/
345+
void send(Set<DataWithMediaType> items) throws IOException;
346+
307347
void complete();
308348

309349
void completeWithError(Throwable failure);

‎spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterReturnValueHandler.java

Lines changed: 10 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
import java.nio.charset.StandardCharsets;
2121
import java.util.ArrayList;
2222
import java.util.List;
23+
import java.util.Set;
2324
import java.util.function.Consumer;
2425

2526
import jakarta.servlet.ServletRequest;
@@ -202,14 +203,22 @@ public HttpMessageConvertingHandler(ServerHttpResponse outputMessage, DeferredRe
202203
@Override
203204
public void send(Object data, @Nullable MediaType mediaType) throws IOException {
204205
sendInternal(data, mediaType);
206+
this.outputMessage.flush();
207+
}
208+
209+
@Override
210+
public void send(Set<ResponseBodyEmitter.DataWithMediaType> items) throws IOException {
211+
for (ResponseBodyEmitter.DataWithMediaType item : items) {
212+
sendInternal(item.getData(), item.getMediaType());
213+
}
214+
this.outputMessage.flush();
205215
}
206216

207217
@SuppressWarnings("unchecked")
208218
private <T> void sendInternal(T data, @Nullable MediaType mediaType) throws IOException {
209219
for (HttpMessageConverter<?> converter : ResponseBodyEmitterReturnValueHandler.this.sseMessageConverters) {
210220
if (converter.canWrite(data.getClass(), mediaType)) {
211221
((HttpMessageConverter<T>) converter).write(data, mediaType, this.outputMessage);
212-
this.outputMessage.flush();
213222
return;
214223
}
215224
}

‎spring-webmvc/src/main/java/org/springframework/web/servlet/mvc/method/annotation/SseEmitter.java

Lines changed: 2 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
/*
2-
* Copyright 2002-2021 the original author or authors.
2+
* Copyright 2002-2023 the original author or authors.
33
*
44
* Licensed under the Apache License, Version 2.0 (the "License");
55
* you may not use this file except in compliance with the License.
@@ -123,9 +123,7 @@ public void send(Object object, @Nullable MediaType mediaType) throws IOExceptio
123123
public void send(SseEventBuilder builder) throws IOException {
124124
Set<DataWithMediaType> dataToSend = builder.build();
125125
synchronized (this) {
126-
for (DataWithMediaType entry : dataToSend) {
127-
super.send(entry.getData(), entry.getMediaType());
128-
}
126+
super.send(dataToSend);
129127
}
130128
}
131129

‎spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ReactiveTypeHandlerTests.java

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -365,6 +365,11 @@ public void send(Object data, MediaType mediaType) throws IOException {
365365
this.values.add(data);
366366
}
367367

368+
@Override
369+
public void send(Set<ResponseBodyEmitter.DataWithMediaType> items) throws IOException {
370+
items.forEach(item -> this.values.add(item.getData()));
371+
}
372+
368373
@Override
369374
public void complete() {
370375
}

‎spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/ResponseBodyEmitterTests.java

Lines changed: 15 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -30,9 +30,9 @@
3030
import static org.assertj.core.api.Assertions.assertThatIOException;
3131
import static org.assertj.core.api.Assertions.assertThatIllegalStateException;
3232
import static org.mockito.ArgumentMatchers.any;
33+
import static org.mockito.ArgumentMatchers.anySet;
3334
import static org.mockito.BDDMockito.willThrow;
3435
import static org.mockito.Mockito.mock;
35-
import static org.mockito.Mockito.times;
3636
import static org.mockito.Mockito.verify;
3737
import static org.mockito.Mockito.verifyNoMoreInteractions;
3838

@@ -52,56 +52,54 @@ public class ResponseBodyEmitterTests {
5252

5353

5454
@Test
55-
public void sendBeforeHandlerInitialized() throws Exception {
55+
void sendBeforeHandlerInitialized() throws Exception {
5656
this.emitter.send("foo", MediaType.TEXT_PLAIN);
5757
this.emitter.send("bar", MediaType.TEXT_PLAIN);
5858
this.emitter.complete();
5959
verifyNoMoreInteractions(this.handler);
6060

6161
this.emitter.initialize(this.handler);
62-
verify(this.handler).send("foo", MediaType.TEXT_PLAIN);
63-
verify(this.handler).send("bar", MediaType.TEXT_PLAIN);
62+
verify(this.handler).send(anySet());
6463
verify(this.handler).complete();
6564
verifyNoMoreInteractions(this.handler);
6665
}
6766

6867
@Test
69-
public void sendDuplicateBeforeHandlerInitialized() throws Exception {
68+
void sendDuplicateBeforeHandlerInitialized() throws Exception {
7069
this.emitter.send("foo", MediaType.TEXT_PLAIN);
7170
this.emitter.send("foo", MediaType.TEXT_PLAIN);
7271
this.emitter.complete();
7372
verifyNoMoreInteractions(this.handler);
7473

7574
this.emitter.initialize(this.handler);
76-
verify(this.handler, times(2)).send("foo", MediaType.TEXT_PLAIN);
75+
verify(this.handler).send(anySet());
7776
verify(this.handler).complete();
7877
verifyNoMoreInteractions(this.handler);
7978
}
8079

8180
@Test
82-
public void sendBeforeHandlerInitializedWithError() throws Exception {
81+
void sendBeforeHandlerInitializedWithError() throws Exception {
8382
IllegalStateException ex = new IllegalStateException();
8483
this.emitter.send("foo", MediaType.TEXT_PLAIN);
8584
this.emitter.send("bar", MediaType.TEXT_PLAIN);
8685
this.emitter.completeWithError(ex);
8786
verifyNoMoreInteractions(this.handler);
8887

8988
this.emitter.initialize(this.handler);
90-
verify(this.handler).send("foo", MediaType.TEXT_PLAIN);
91-
verify(this.handler).send("bar", MediaType.TEXT_PLAIN);
89+
verify(this.handler).send(anySet());
9290
verify(this.handler).completeWithError(ex);
9391
verifyNoMoreInteractions(this.handler);
9492
}
9593

9694
@Test
97-
public void sendFailsAfterComplete() throws Exception {
95+
void sendFailsAfterComplete() throws Exception {
9896
this.emitter.complete();
9997
assertThatIllegalStateException().isThrownBy(() ->
10098
this.emitter.send("foo"));
10199
}
102100

103101
@Test
104-
public void sendAfterHandlerInitialized() throws Exception {
102+
void sendAfterHandlerInitialized() throws Exception {
105103
this.emitter.initialize(this.handler);
106104
verify(this.handler).onTimeout(any());
107105
verify(this.handler).onError(any());
@@ -119,7 +117,7 @@ public void sendAfterHandlerInitialized() throws Exception {
119117
}
120118

121119
@Test
122-
public void sendAfterHandlerInitializedWithError() throws Exception {
120+
void sendAfterHandlerInitializedWithError() throws Exception {
123121
this.emitter.initialize(this.handler);
124122
verify(this.handler).onTimeout(any());
125123
verify(this.handler).onError(any());
@@ -138,7 +136,7 @@ public void sendAfterHandlerInitializedWithError() throws Exception {
138136
}
139137

140138
@Test
141-
public void sendWithError() throws Exception {
139+
void sendWithError() throws Exception {
142140
this.emitter.initialize(this.handler);
143141
verify(this.handler).onTimeout(any());
144142
verify(this.handler).onError(any());
@@ -154,7 +152,7 @@ public void sendWithError() throws Exception {
154152
}
155153

156154
@Test
157-
public void onTimeoutBeforeHandlerInitialized() throws Exception {
155+
void onTimeoutBeforeHandlerInitialized() throws Exception {
158156
Runnable runnable = mock();
159157
this.emitter.onTimeout(runnable);
160158
this.emitter.initialize(this.handler);
@@ -169,7 +167,7 @@ public void onTimeoutBeforeHandlerInitialized() throws Exception {
169167
}
170168

171169
@Test
172-
public void onTimeoutAfterHandlerInitialized() throws Exception {
170+
void onTimeoutAfterHandlerInitialized() throws Exception {
173171
this.emitter.initialize(this.handler);
174172

175173
ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);
@@ -185,7 +183,7 @@ public void onTimeoutAfterHandlerInitialized() throws Exception {
185183
}
186184

187185
@Test
188-
public void onCompletionBeforeHandlerInitialized() throws Exception {
186+
void onCompletionBeforeHandlerInitialized() throws Exception {
189187
Runnable runnable = mock();
190188
this.emitter.onCompletion(runnable);
191189
this.emitter.initialize(this.handler);
@@ -200,7 +198,7 @@ public void onCompletionBeforeHandlerInitialized() throws Exception {
200198
}
201199

202200
@Test
203-
public void onCompletionAfterHandlerInitialized() throws Exception {
201+
void onCompletionAfterHandlerInitialized() throws Exception {
204202
this.emitter.initialize(this.handler);
205203

206204
ArgumentCaptor<Runnable> captor = ArgumentCaptor.forClass(Runnable.class);

‎spring-webmvc/src/test/java/org/springframework/web/servlet/mvc/method/annotation/SseEmitterTests.java

Lines changed: 28 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,12 +20,14 @@
2020
import java.nio.charset.StandardCharsets;
2121
import java.util.ArrayList;
2222
import java.util.List;
23+
import java.util.Set;
2324
import java.util.function.Consumer;
2425

2526
import org.junit.jupiter.api.BeforeEach;
2627
import org.junit.jupiter.api.Test;
2728

2829
import org.springframework.http.MediaType;
30+
import org.springframework.lang.Nullable;
2931

3032
import static org.assertj.core.api.Assertions.assertThat;
3133
import static org.springframework.web.servlet.mvc.method.annotation.SseEmitter.event;
@@ -60,6 +62,7 @@ public void send() throws Exception {
6062
this.handler.assertObject(0, "data:", TEXT_PLAIN_UTF8);
6163
this.handler.assertObject(1, "foo");
6264
this.handler.assertObject(2, "\n\n", TEXT_PLAIN_UTF8);
65+
this.handler.assertWriteCount(1);
6366
}
6467

6568
@Test
@@ -69,12 +72,14 @@ public void sendWithMediaType() throws Exception {
6972
this.handler.assertObject(0, "data:", TEXT_PLAIN_UTF8);
7073
this.handler.assertObject(1, "foo", MediaType.TEXT_PLAIN);
7174
this.handler.assertObject(2, "\n\n", TEXT_PLAIN_UTF8);
75+
this.handler.assertWriteCount(1);
7276
}
7377

7478
@Test
7579
public void sendEventEmpty() throws Exception {
7680
this.emitter.send(event());
7781
this.handler.assertSentObjectCount(0);
82+
this.handler.assertWriteCount(0);
7883
}
7984

8085
@Test
@@ -84,6 +89,7 @@ public void sendEventWithDataLine() throws Exception {
8489
this.handler.assertObject(0, "data:", TEXT_PLAIN_UTF8);
8590
this.handler.assertObject(1, "foo");
8691
this.handler.assertObject(2, "\n\n", TEXT_PLAIN_UTF8);
92+
this.handler.assertWriteCount(1);
8793
}
8894

8995
@Test
@@ -95,6 +101,7 @@ public void sendEventWithTwoDataLines() throws Exception {
95101
this.handler.assertObject(2, "\ndata:", TEXT_PLAIN_UTF8);
96102
this.handler.assertObject(3, "bar");
97103
this.handler.assertObject(4, "\n\n", TEXT_PLAIN_UTF8);
104+
this.handler.assertWriteCount(1);
98105
}
99106

100107
@Test
@@ -104,6 +111,7 @@ public void sendEventFull() throws Exception {
104111
this.handler.assertObject(0, ":blah\nevent:test\nretry:5000\nid:1\ndata:", TEXT_PLAIN_UTF8);
105112
this.handler.assertObject(1, "foo");
106113
this.handler.assertObject(2, "\n\n", TEXT_PLAIN_UTF8);
114+
this.handler.assertWriteCount(1);
107115
}
108116

109117
@Test
@@ -115,14 +123,17 @@ public void sendEventFullWithTwoDataLinesInTheMiddle() throws Exception {
115123
this.handler.assertObject(2, "\ndata:", TEXT_PLAIN_UTF8);
116124
this.handler.assertObject(3, "bar");
117125
this.handler.assertObject(4, "\nevent:test\nretry:5000\nid:1\n\n", TEXT_PLAIN_UTF8);
126+
this.handler.assertWriteCount(1);
118127
}
119128

120129

121130
private static class TestHandler implements ResponseBodyEmitter.Handler {
122131

123-
private List<Object> objects = new ArrayList<>();
132+
private final List<Object> objects = new ArrayList<>();
124133

125-
private List<MediaType> mediaTypes = new ArrayList<>();
134+
private final List<MediaType> mediaTypes = new ArrayList<>();
135+
136+
private int writeCount;
126137

127138

128139
public void assertSentObjectCount(int size) {
@@ -139,10 +150,24 @@ public void assertObject(int index, Object object, MediaType mediaType) {
139150
assertThat(this.mediaTypes.get(index)).isEqualTo(mediaType);
140151
}
141152

153+
public void assertWriteCount(int writeCount) {
154+
assertThat(this.writeCount).isEqualTo(writeCount);
155+
}
156+
142157
@Override
143-
public void send(Object data, MediaType mediaType) throws IOException {
158+
public void send(Object data, @Nullable MediaType mediaType) throws IOException {
144159
this.objects.add(data);
145160
this.mediaTypes.add(mediaType);
161+
this.writeCount++;
162+
}
163+
164+
@Override
165+
public void send(Set<ResponseBodyEmitter.DataWithMediaType> items) throws IOException {
166+
for (ResponseBodyEmitter.DataWithMediaType item : items) {
167+
this.objects.add(item.getData());
168+
this.mediaTypes.add(item.getMediaType());
169+
}
170+
this.writeCount++;
146171
}
147172

148173
@Override

0 commit comments

Comments
 (0)