001/*
002 * Copyright 2002-2020 the original author or authors.
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License");
005 * you may not use this file except in compliance with the License.
006 * You may obtain a copy of the License at
007 *
008 *      https://www.apache.org/licenses/LICENSE-2.0
009 *
010 * Unless required by applicable law or agreed to in writing, software
011 * distributed under the License is distributed on an "AS IS" BASIS,
012 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
013 * See the License for the specific language governing permissions and
014 * limitations under the License.
015 */
016
017package org.springframework.web.servlet.mvc.method.annotation;
018
019import java.io.IOException;
020import java.nio.charset.StandardCharsets;
021import java.util.Collections;
022import java.util.LinkedHashSet;
023import java.util.Set;
024
025import org.springframework.http.HttpHeaders;
026import org.springframework.http.MediaType;
027import org.springframework.http.server.ServerHttpResponse;
028import org.springframework.lang.Nullable;
029import org.springframework.util.ObjectUtils;
030import org.springframework.util.StringUtils;
031
032/**
033 * A specialization of {@link ResponseBodyEmitter} for sending
034 * <a href="https://www.w3.org/TR/eventsource/">Server-Sent Events</a>.
035 *
036 * @author Rossen Stoyanchev
037 * @author Juergen Hoeller
038 * @since 4.2
039 */
040public class SseEmitter extends ResponseBodyEmitter {
041
042        private static final MediaType TEXT_PLAIN = new MediaType("text", "plain", StandardCharsets.UTF_8);
043
044        /**
045         * Create a new SseEmitter instance.
046         */
047        public SseEmitter() {
048                super();
049        }
050
051        /**
052         * Create a SseEmitter with a custom timeout value.
053         * <p>By default not set in which case the default configured in the MVC
054         * Java Config or the MVC namespace is used, or if that's not set, then the
055         * timeout depends on the default of the underlying server.
056         * @param timeout the timeout value in milliseconds
057         * @since 4.2.2
058         */
059        public SseEmitter(Long timeout) {
060                super(timeout);
061        }
062
063
064        @Override
065        protected void extendResponse(ServerHttpResponse outputMessage) {
066                super.extendResponse(outputMessage);
067
068                HttpHeaders headers = outputMessage.getHeaders();
069                if (headers.getContentType() == null) {
070                        headers.setContentType(MediaType.TEXT_EVENT_STREAM);
071                }
072        }
073
074        /**
075         * Send the object formatted as a single SSE "data" line. It's equivalent to:
076         * <pre>
077         * // static import of SseEmitter.*
078         *
079         * SseEmitter emitter = new SseEmitter();
080         * emitter.send(event().data(myObject));
081         * </pre>
082         * <p>Please, see {@link ResponseBodyEmitter#send(Object) parent Javadoc}
083         * for important notes on exception handling.
084         * @param object the object to write
085         * @throws IOException raised when an I/O error occurs
086         * @throws java.lang.IllegalStateException wraps any other errors
087         */
088        @Override
089        public void send(Object object) throws IOException {
090                send(object, null);
091        }
092
093        /**
094         * Send the object formatted as a single SSE "data" line. It's equivalent to:
095         * <pre>
096         * // static import of SseEmitter.*
097         *
098         * SseEmitter emitter = new SseEmitter();
099         * emitter.send(event().data(myObject, MediaType.APPLICATION_JSON));
100         * </pre>
101         * <p>Please, see {@link ResponseBodyEmitter#send(Object) parent Javadoc}
102         * for important notes on exception handling.
103         * @param object the object to write
104         * @param mediaType a MediaType hint for selecting an HttpMessageConverter
105         * @throws IOException raised when an I/O error occurs
106         */
107        @Override
108        public void send(Object object, @Nullable MediaType mediaType) throws IOException {
109                send(event().data(object, mediaType));
110        }
111
112        /**
113         * Send an SSE event prepared with the given builder. For example:
114         * <pre>
115         * // static import of SseEmitter
116         * SseEmitter emitter = new SseEmitter();
117         * emitter.send(event().name("update").id("1").data(myObject));
118         * </pre>
119         * @param builder a builder for an SSE formatted event.
120         * @throws IOException raised when an I/O error occurs
121         */
122        public void send(SseEventBuilder builder) throws IOException {
123                Set<DataWithMediaType> dataToSend = builder.build();
124                synchronized (this) {
125                        for (DataWithMediaType entry : dataToSend) {
126                                super.send(entry.getData(), entry.getMediaType());
127                        }
128                }
129        }
130
131        @Override
132        public String toString() {
133                return "SseEmitter@" + ObjectUtils.getIdentityHexString(this);
134        }
135
136
137        public static SseEventBuilder event() {
138                return new SseEventBuilderImpl();
139        }
140
141
142        /**
143         * A builder for an SSE event.
144         */
145        public interface SseEventBuilder {
146
147                /**
148                 * Add an SSE "id" line.
149                 */
150                SseEventBuilder id(String id);
151
152                /**
153                 * Add an SSE "event" line.
154                 */
155                SseEventBuilder name(String eventName);
156
157                /**
158                 * Add an SSE "retry" line.
159                 */
160                SseEventBuilder reconnectTime(long reconnectTimeMillis);
161
162                /**
163                 * Add an SSE "comment" line.
164                 */
165                SseEventBuilder comment(String comment);
166
167                /**
168                 * Add an SSE "data" line.
169                 */
170                SseEventBuilder data(Object object);
171
172                /**
173                 * Add an SSE "data" line.
174                 */
175                SseEventBuilder data(Object object, @Nullable MediaType mediaType);
176
177                /**
178                 * Return one or more Object-MediaType  pairs to write via
179                 * {@link #send(Object, MediaType)}.
180                 * @since 4.2.3
181                 */
182                Set<DataWithMediaType> build();
183        }
184
185
186        /**
187         * Default implementation of SseEventBuilder.
188         */
189        private static class SseEventBuilderImpl implements SseEventBuilder {
190
191                private final Set<DataWithMediaType> dataToSend = new LinkedHashSet<>(4);
192
193                @Nullable
194                private StringBuilder sb;
195
196                @Override
197                public SseEventBuilder id(String id) {
198                        append("id:").append(id).append("\n");
199                        return this;
200                }
201
202                @Override
203                public SseEventBuilder name(String name) {
204                        append("event:").append(name).append("\n");
205                        return this;
206                }
207
208                @Override
209                public SseEventBuilder reconnectTime(long reconnectTimeMillis) {
210                        append("retry:").append(String.valueOf(reconnectTimeMillis)).append("\n");
211                        return this;
212                }
213
214                @Override
215                public SseEventBuilder comment(String comment) {
216                        append(":").append(comment).append("\n");
217                        return this;
218                }
219
220                @Override
221                public SseEventBuilder data(Object object) {
222                        return data(object, null);
223                }
224
225                @Override
226                public SseEventBuilder data(Object object, @Nullable MediaType mediaType) {
227                        append("data:");
228                        saveAppendedText();
229                        this.dataToSend.add(new DataWithMediaType(object, mediaType));
230                        append("\n");
231                        return this;
232                }
233
234                SseEventBuilderImpl append(String text) {
235                        if (this.sb == null) {
236                                this.sb = new StringBuilder();
237                        }
238                        this.sb.append(text);
239                        return this;
240                }
241
242                @Override
243                public Set<DataWithMediaType> build() {
244                        if (!StringUtils.hasLength(this.sb) && this.dataToSend.isEmpty()) {
245                                return Collections.emptySet();
246                        }
247                        append("\n");
248                        saveAppendedText();
249                        return this.dataToSend;
250                }
251
252                private void saveAppendedText() {
253                        if (this.sb != null) {
254                                this.dataToSend.add(new DataWithMediaType(this.sb.toString(), TEXT_PLAIN));
255                                this.sb = null;
256                        }
257                }
258        }
259
260}