001/*
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.juneau.rest.server.sse;
018
019
020import java.io.*;
021import java.time.*;
022import java.util.concurrent.*;
023
024import org.apache.juneau.marshall.marshaller.*;
025import org.apache.juneau.marshall.sse.*;
026import org.apache.juneau.rest.server.*;
027import org.apache.juneau.rest.server.util.*;
028import static org.apache.juneau.commons.utils.Shorts.*;
029
030/**
031 * Fluent SSE response helper.
032 */
033@SuppressWarnings({
034   "resource" // The heartbeat and negotiated writer are owned by this object/response and are closed in close().
035})
036public class SseResponseSupport implements AutoCloseable {
037
038   private final RestResponse response;
039   private final FinishablePrintWriter writer;
040   private final ScheduledExecutorService scheduler;
041   private SseHeartbeat heartbeat;
042
043   /**
044    * Constructor.
045    *
046    * @param response The REST response. Must not be <jk>null</jk>.
047    * @throws IOException If the writer could not be created.
048    */
049   public SseResponseSupport(RestResponse response) throws IOException {
050      this.response = reqnn("response", response);
051      this.scheduler = response.getContext().getBeanStore().getBean(ScheduledExecutorService.class).orElse(null);
052      response.setContentType("text/event-stream");
053      response.setHeader("Cache-Control", "no-cache");
054      response.setHeader("X-Accel-Buffering", "no");
055      response.setHeader("X-Content-Type-Options", "nosniff");
056      response.setHeader("Content-Encoding", "identity");
057      writer = response.getNegotiatedWriter();
058   }
059
060   /**
061    * Starts periodic heartbeat comments.
062    *
063    * @param interval The heartbeat interval.
064    * @return This object.
065    */
066   public SseResponseSupport heartbeat(Duration interval) {
067      if (scheduler != null) {
068         if (heartbeat != null)
069            heartbeat.close();
070         heartbeat = SseHeartbeat.start(scheduler, writer, interval);
071      }
072      return this;
073   }
074
075   /**
076    * Sends an SSE event.
077    *
078    * @param event The event.
079    * @return This object.
080    * @throws IOException If an I/O error occurred.
081    */
082   public SseResponseSupport sendEvent(SseEvent event) throws IOException {
083      Sse.DEFAULT.write(event, writer);
084      return this;
085   }
086
087   /**
088    * Sends an SSE event from name+data values.
089    *
090    * @param name The event name. Can be <jk>null</jk> (dispatches as the default event type).
091    * @param data The event data. Can be <jk>null</jk> (no data lines emitted).
092    * @return This object.
093    * @throws IOException If an I/O error occurred.
094    */
095   public SseResponseSupport sendEvent(String name, Object data) throws IOException {
096      return sendEvent(new SseEvent(name, data == null ? null : data.toString()));
097   }
098
099   /**
100    * Sends a heartbeat/comment line.
101    *
102    * @param value The comment value. Can be <jk>null</jk> (treated as an empty comment).
103    * @return This object.
104    * @throws IOException If an I/O error occurred.
105    */
106   public SseResponseSupport comment(String value) throws IOException {
107      SseSerializer.writeComment(writer, value);
108      return this;
109   }
110
111   /**
112    * Flushes pending output.
113    *
114    * @return This object.
115    * @throws IOException If an I/O error occurred.
116    */
117   public SseResponseSupport flush() throws IOException {
118      writer.flush();
119      response.flushBuffer();
120      return this;
121   }
122
123   /**
124    * Drains a subscription until disconnect or interruption.
125    *
126    * @param subscription The subscription. Must not be <jk>null</jk>.
127    * @return This object.
128    * @throws IOException If an I/O error occurred.
129    */
130   public SseResponseSupport sendFrom(SseSubscription subscription) throws IOException {
131      reqnn("subscription", subscription);
132      try {
133         while (! subscription.isClosed()) {
134            sendEvent(subscription.take());
135            flush();
136         }
137      } catch (InterruptedException e) {
138         Thread.currentThread().interrupt();
139      } finally {
140         subscription.close();
141      }
142      return this;
143   }
144
145   @Override /* AutoCloseable */
146   public void close() {
147      if (heartbeat != null)
148         heartbeat.close();
149   }
150}