001/* 002 * (C) Copyright 2011 Nuxeo SA (http://nuxeo.com/) and others. 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 * http://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 * Contributors: 017 * Florent Guillaume 018 */ 019package org.nuxeo.ecm.core.io.download; 020 021import java.io.BufferedOutputStream; 022import java.io.ByteArrayOutputStream; 023import java.io.File; 024import java.io.FileInputStream; 025import java.io.FileOutputStream; 026import java.io.IOException; 027import java.io.OutputStream; 028import java.io.OutputStreamWriter; 029import java.io.PrintWriter; 030 031import javax.servlet.ServletOutputStream; 032 033import org.apache.commons.io.IOUtils; 034import org.apache.commons.logging.Log; 035import org.apache.commons.logging.LogFactory; 036 037import org.nuxeo.runtime.api.Framework; 038 039/** 040 * A {@link ServletOutputStream} that buffers everything until {@link #stopBuffering()} is called. 041 * <p> 042 * There may only be one such instance per thread. 043 * <p> 044 * Buffering is done first in memory, then on disk if the size exceeds a limit. 045 */ 046public class BufferingServletOutputStream extends ServletOutputStream { 047 048 private static final Log log = LogFactory.getLog(BufferingServletOutputStream.class); 049 050 /** Initial memory buffer size. */ 051 public static final int INITIAL = 4 * 1024; // 4 KB 052 053 /** Maximum memory buffer size, after this a file is used. */ 054 public static final int MAX = 64 * 1024; // 64 KB 055 056 /** Used for 0-length writes. */ 057 private final static OutputStream EMPTY = new ByteArrayOutputStream(0); 058 059 /** Have we stopped buffering to pass writes directly to the output stream. */ 060 protected boolean streaming; 061 062 protected boolean needsFlush; 063 064 protected boolean needsClose; 065 066 protected final OutputStream outputStream; 067 068 protected PrintWriter writer; 069 070 protected ByteArrayOutputStream memory; 071 072 protected OutputStream file; 073 074 protected File tmp; 075 076 /** 077 * A {@link ServletOutputStream} wrapper that buffers everything until {@link #stopBuffering()} is called. 078 * <p> 079 * {@link #stopBuffering()} <b>MUST</b> be called in a {@code finally} statement in order for resources to be closed 080 * properly. 081 * 082 * @param outputStream the underlying output stream 083 */ 084 public BufferingServletOutputStream(OutputStream outputStream) { 085 this.outputStream = outputStream; 086 } 087 088 public PrintWriter getWriter() { 089 if (writer == null) { 090 writer = new PrintWriter(new OutputStreamWriter(this)); 091 } 092 return writer; 093 } 094 095 /** 096 * Finds the proper output stream where we can write {@code len} bytes. 097 */ 098 protected OutputStream getOutputStream(int len) throws IOException { 099 if (streaming) { 100 return outputStream; 101 } 102 if (len == 0) { 103 return EMPTY; 104 } 105 if (file != null) { 106 // already to file 107 return file; 108 } 109 int total; 110 if (memory == null) { 111 // no buffer yet 112 if (len <= MAX) { 113 memory = new ByteArrayOutputStream(Math.max(INITIAL, len)); 114 return memory; 115 } 116 total = len; 117 } else { 118 total = memory.size() + len; 119 } 120 if (total <= MAX) { 121 return memory; 122 } else { 123 // switch to a file 124 createTempFile(); 125 file = new BufferedOutputStream(new FileOutputStream(tmp)); 126 if (memory != null) { 127 memory.writeTo(file); 128 memory = null; 129 } 130 return file; 131 } 132 } 133 134 protected void createTempFile() throws IOException { 135 tmp = Framework.createTempFile("nxout", null); 136 } 137 138 @Override 139 public void write(int b) throws IOException { 140 getOutputStream(1).write(b); 141 } 142 143 @Override 144 public void write(byte[] b) throws IOException { 145 getOutputStream(b.length).write(b); 146 } 147 148 @Override 149 public void write(byte[] b, int off, int len) throws IOException { 150 getOutputStream(len).write(b, off, len); 151 } 152 153 /** 154 * This implementation does nothing, we still want to keep buffering and not flush. 155 * <p> 156 * {@inheritDoc} 157 */ 158 @Override 159 public void flush() throws IOException { 160 if (streaming) { 161 outputStream.flush(); 162 } else { 163 needsFlush = true; 164 } 165 } 166 167 /** 168 * This implementation does nothing, we still want to keep the buffer until {@link #stopBuffering()} time. 169 * <p> 170 * {@inheritDoc} 171 */ 172 @Override 173 public void close() throws IOException { 174 if (streaming) { 175 outputStream.close(); 176 } else { 177 needsClose = true; 178 } 179 180 } 181 182 /** 183 * Writes any buffered data to the underlying {@link OutputStream} and from now on don't buffer anymore. 184 */ 185 public void stopBuffering() throws IOException { 186 if (streaming) { 187 return; 188 } 189 if (writer != null) { 190 writer.flush(); // don't close, streaming needs it 191 } 192 streaming = true; 193 if (log.isDebugEnabled()) { 194 long len; 195 if (memory != null) { 196 len = memory.size(); 197 } else if (file != null) { 198 len = tmp.length(); 199 } else { 200 len = 0; 201 } 202 log.debug("buffered bytes: " + len); 203 } 204 boolean clientAbort = false; 205 try { 206 if (memory != null) { 207 memory.writeTo(outputStream); 208 } else if (file != null) { 209 try { 210 try { 211 file.flush(); 212 } finally { 213 file.close(); 214 } 215 FileInputStream in = new FileInputStream(tmp); 216 try { 217 IOUtils.copy(in, outputStream); 218 } catch (IOException e) { 219 if (DownloadHelper.isClientAbortError(e)) { 220 DownloadHelper.logClientAbort(e); 221 clientAbort = true; 222 } else { 223 throw e; 224 } 225 } finally { 226 in.close(); 227 } 228 } finally { 229 tmp.delete(); 230 } 231 } 232 } catch (IOException e) { 233 if (DownloadHelper.isClientAbortError(e)) { 234 if (!clientAbort) { 235 DownloadHelper.logClientAbort(e); 236 clientAbort = true; 237 } 238 } else { 239 throw e; 240 } 241 } finally { 242 memory = null; 243 file = null; 244 tmp = null; 245 try { 246 if (needsFlush) { 247 outputStream.flush(); 248 } 249 } catch (IOException e) { 250 if (DownloadHelper.isClientAbortError(e)) { 251 if (!clientAbort) { 252 DownloadHelper.logClientAbort(e); 253 } 254 } else { 255 throw e; 256 } 257 } finally { 258 if (needsClose) { 259 outputStream.close(); 260 } 261 } 262 } 263 } 264 265 /** 266 * Tells the given {@link OutputStream} to stop buffering (if it was). 267 */ 268 public static void stopBuffering(OutputStream out) throws IOException { 269 if (out instanceof BufferingServletOutputStream) { 270 ((BufferingServletOutputStream) out).stopBuffering(); 271 } 272 } 273 274}