001/*
002 * Copyright 2002-2019 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.Charset;
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.util.ObjectUtils;
029import org.springframework.util.StringUtils;
030
031/**
032 * A specialization of {@link ResponseBodyEmitter} for sending
033 * <a href="https://www.w3.org/TR/eventsource/">Server-Sent Events</a>.
034 *
035 * @author Rossen Stoyanchev
036 * @author Juergen Hoeller
037 * @since 4.2
038 */
039public class SseEmitter extends ResponseBodyEmitter {
040
041        static final MediaType TEXT_PLAIN = new MediaType("text", "plain", Charset.forName("UTF-8"));
042
043        static final MediaType TEXT_EVENTSTREAM = new MediaType("text", "event-stream", Charset.forName("UTF-8"));
044
045
046        /**
047         * Create a new SseEmitter instance.
048         */
049        public SseEmitter() {
050                super();
051        }
052
053        /**
054         * Create a SseEmitter with a custom timeout value.
055         * <p>By default not set in which case the default configured in the MVC
056         * Java Config or the MVC namespace is used, or if that's not set, then the
057         * timeout depends on the default of the underlying server.
058         * @param timeout timeout value in milliseconds
059         * @since 4.2.2
060         */
061        public SseEmitter(Long timeout) {
062                super(timeout);
063        }
064
065
066        @Override
067        protected void extendResponse(ServerHttpResponse outputMessage) {
068                super.extendResponse(outputMessage);
069
070                HttpHeaders headers = outputMessage.getHeaders();
071                if (headers.getContentType() == null) {
072                        headers.setContentType(TEXT_EVENTSTREAM);
073                }
074        }
075
076        /**
077         * Send the object formatted as a single SSE "data" line. It's equivalent to:
078         * <pre>
079         * // static import of SseEmitter.*
080         *
081         * SseEmitter emitter = new SseEmitter();
082         * emitter.send(event().data(myObject));
083         * </pre>
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         * @param object the object to write
102         * @param mediaType a MediaType hint for selecting an HttpMessageConverter
103         * @throws IOException raised when an I/O error occurs
104         */
105        @Override
106        public void send(Object object, MediaType mediaType) throws IOException {
107                if (object != null) {
108                        send(event().data(object, mediaType));
109                }
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, 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<DataWithMediaType>(4);
192
193                private StringBuilder sb;
194
195                @Override
196                public SseEventBuilder id(String id) {
197                        append("id:").append(id != null ? id : "").append("\n");
198                        return this;
199                }
200
201                @Override
202                public SseEventBuilder name(String name) {
203                        append("event:").append(name != null ? name : "").append("\n");
204                        return this;
205                }
206
207                @Override
208                public SseEventBuilder reconnectTime(long reconnectTimeMillis) {
209                        append("retry:").append(String.valueOf(reconnectTimeMillis)).append("\n");
210                        return this;
211                }
212
213                @Override
214                public SseEventBuilder comment(String comment) {
215                        append(":").append(comment != null ? comment : "").append("\n");
216                        return this;
217                }
218
219                @Override
220                public SseEventBuilder data(Object object) {
221                        return data(object, null);
222                }
223
224                @Override
225                public SseEventBuilder data(Object object, MediaType mediaType) {
226                        append("data:");
227                        saveAppendedText();
228                        this.dataToSend.add(new DataWithMediaType(object, mediaType));
229                        append("\n");
230                        return this;
231                }
232
233                SseEventBuilderImpl append(String text) {
234                        if (this.sb == null) {
235                                this.sb = new StringBuilder();
236                        }
237                        this.sb.append(text);
238                        return this;
239                }
240
241                @Override
242                public Set<DataWithMediaType> build() {
243                        if (!StringUtils.hasLength(this.sb) && this.dataToSend.isEmpty()) {
244                                return Collections.<DataWithMediaType>emptySet();
245                        }
246                        append("\n");
247                        saveAppendedText();
248                        return this.dataToSend;
249                }
250
251                private void saveAppendedText() {
252                        if (this.sb != null) {
253                                this.dataToSend.add(new DataWithMediaType(this.sb.toString(), TEXT_PLAIN));
254                                this.sb = null;
255                        }
256                }
257        }
258
259}