org.spearce.jgit/src/org/spearce/jgit/util/TemporaryBuffer.java

   1 /*
   2  * Copyright (C) 2008, Shawn O. Pearce <spearce@spearce.org>
   3  *
   4  * All rights reserved.
   5  *
   6  * Redistribution and use in source and binary forms, with or
   7  * without modification, are permitted provided that the following
   8  * conditions are met:
   9  *
  10  * - Redistributions of source code must retain the above copyright
  11  *   notice, this list of conditions and the following disclaimer.
  12  *
  13  * - Redistributions in binary form must reproduce the above
  14  *   copyright notice, this list of conditions and the following
  15  *   disclaimer in the documentation and/or other materials provided
  16  *   with the distribution.
  17  *
  18  * - Neither the name of the Git Development Community nor the
  19  *   names of its contributors may be used to endorse or promote
  20  *   products derived from this software without specific prior
  21  *   written permission.
  22  *
  23  * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND
  24  * CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
  25  * INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
  26  * OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
  27  * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
  28  * CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
  29  * SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
  30  * NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
  31  * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
  32  * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
  33  * STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
  34  * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
  35  * ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
  36  */
  37
  38 package org.spearce.jgit.util;
  39
  40 import java.io.BufferedOutputStream;
  41 import java.io.File;
  42 import java.io.FileInputStream;
  43 import java.io.FileOutputStream;
  44 import java.io.IOException;
  45 import java.io.InputStream;
  46 import java.io.OutputStream;
  47 import java.util.ArrayList;
  48
  49 import org.spearce.jgit.lib.NullProgressMonitor;
  50 import org.spearce.jgit.lib.ProgressMonitor;
  51
  52 /**
  53  * A fully buffered output stream using local disk storage for large data.
  54  * <p>
  55  * Initially this output stream buffers to memory, like ByteArrayOutputStream
  56  * might do, but it shifts to using an on disk temporary file if the output gets
  57  * too large.
  58  * <p>
  59  * The content of this buffered stream may be sent to another OutputStream only
  60  * after this stream has been properly closed by {@link #close()}.
  61  */
  62 public class TemporaryBuffer extends OutputStream {
  63         static final int DEFAULT_IN_CORE_LIMIT = 1024 * 1024;
  64
  65         /** Chain of data, if we are still completely in-core; otherwise null. */
  66         private ArrayList<Block> blocks;
  67
  68         /**
  69          * Maximum number of bytes we will permit storing in memory.
  70          * <p>
  71          * When this limit is reached the data will be shifted to a file on disk,
  72          * preventing the JVM heap from growing out of control.
  73          */
  74         private int inCoreLimit;
  75
  76         /**
  77          * Location of our temporary file if we are on disk; otherwise null.
  78          * <p>
  79          * If we exceeded the {@link #inCoreLimit} we nulled out {@link #blocks} and
  80          * created this file instead. All output goes here through {@link #diskOut}.
  81          */
  82         private File onDiskFile;
  83
  84         /** If writing to {@link #onDiskFile} this is a buffered stream to it. */
  85         private OutputStream diskOut;
  86
  87         /** Create a new empty temporary buffer. */
  88         public TemporaryBuffer() {
  89                 inCoreLimit = DEFAULT_IN_CORE_LIMIT;
  90                 blocks = new ArrayList<Block>(inCoreLimit / Block.SZ);
  91                 blocks.add(new Block());
  92         }
  93
  94         @Override
  95         public void write(final int b) throws IOException {
  96                 if (blocks == null) {
  97                         diskOut.write(b);
  98                         return;
  99                 }
 100
 101                 Block s = last();
 102                 if (s.isFull()) {
 103                         if (reachedInCoreLimit()) {
 104                                 diskOut.write(b);
 105                                 return;
 106                         }
 107
 108                         s = new Block();
 109                         blocks.add(s);
 110                 }
 111                 s.buffer[s.count++] = (byte) b;
 112         }
 113
 114         @Override
 115         public void write(final byte[] b, int off, int len) throws IOException {
 116                 if (blocks != null) {
 117                         while (len > 0) {
 118                                 Block s = last();
 119                                 if (s.isFull()) {
 120                                         if (reachedInCoreLimit())
 121                                                 break;
 122
 123                                         s = new Block();
 124                                         blocks.add(s);
 125                                 }
 126
 127                                 final int n = Math.min(Block.SZ - s.count, len);
 128                                 System.arraycopy(b, off, s.buffer, s.count, n);
 129                                 s.count += n;
 130                                 len -= n;
 131                                 off += n;
 132                         }
 133                 }
 134
 135                 if (len > 0)
 136                         diskOut.write(b, off, len);
 137         }
 138
 139         /**
 140          * Copy all bytes remaining on the input stream into this buffer.
 141          *
 142          * @param in
 143          *            the stream to read from, until EOF is reached.
 144          * @throws IOException
 145          *             an error occurred reading from the input stream, or while
 146          *             writing to a local temporary file.
 147          */
 148         public void copy(final InputStream in) throws IOException {
 149                 if (blocks != null) {
 150                         for (;;) {
 151                                 Block s = last();
 152                                 if (s.isFull()) {
 153                                         if (reachedInCoreLimit())
 154                                                 break;
 155                                         s = new Block();
 156                                         blocks.add(s);
 157                                 }
 158
 159                                 final int n = in.read(s.buffer, s.count, Block.SZ - s.count);
 160                                 if (n < 1)
 161                                         return;
 162                                 s.count += n;
 163                         }
 164                 }
 165
 166                 final byte[] tmp = new byte[Block.SZ];
 167                 int n;
 168                 while ((n = in.read(tmp)) > 0)
 169                         diskOut.write(tmp, 0, n);
 170         }
 171
 172         private Block last() {
 173                 return blocks.get(blocks.size() - 1);
 174         }
 175
 176         private boolean reachedInCoreLimit() throws IOException {
 177                 if (blocks.size() * Block.SZ < inCoreLimit)
 178                         return false;
 179
 180                 onDiskFile = File.createTempFile("jgit_", ".buffer");
 181                 diskOut = new FileOutputStream(onDiskFile);
 182
 183                 final Block last = blocks.remove(blocks.size() - 1);
 184                 for (final Block b : blocks)
 185                         diskOut.write(b.buffer, 0, b.count);
 186                 blocks = null;
 187
 188                 diskOut = new BufferedOutputStream(diskOut, Block.SZ);
 189                 diskOut.write(last.buffer, 0, last.count);
 190                 return true;
 191         }
 192
 193         public void close() throws IOException {
 194                 if (diskOut != null) {
 195                         try {
 196                                 diskOut.close();
 197                         } finally {
 198                                 diskOut = null;
 199                         }
 200                 }
 201         }
 202
 203         /**
 204          * Obtain the length (in bytes) of the buffer.
 205          * <p>
 206          * The length is only accurate after {@link #close()} has been invoked.
 207          *
 208          * @return total length of the buffer, in bytes.
 209          */
 210         public long length() {
 211                 if (onDiskFile != null)
 212                         return onDiskFile.length();
 213
 214                 final Block last = last();
 215                 return ((long) blocks.size()) * Block.SZ - (Block.SZ - last.count);
 216         }
 217
 218         /**
 219          * Convert this buffer's contents into a contiguous byte array.
 220          * <p>
 221          * The buffer is only complete after {@link #close()} has been invoked.
 222          *
 223          * @return the complete byte array; length matches {@link #length()}.
 224          * @throws IOException
 225          *             an error occurred reading from a local temporary file
 226          * @throws OutOfMemoryError
 227          *             the buffer cannot fit in memory
 228          */
 229         public byte[] toByteArray() throws IOException {
 230                 final long len = length();
 231                 if (Integer.MAX_VALUE < len)
 232                         throw new OutOfMemoryError("Length exceeds maximum array size");
 233
 234                 final byte[] out = new byte[(int) len];
 235                 if (blocks != null) {
 236                         int outPtr = 0;
 237                         for (final Block b : blocks) {
 238                                 System.arraycopy(b.buffer, 0, out, outPtr, b.count);
 239                                 outPtr += b.count;
 240                         }
 241                 } else {
 242                         final FileInputStream in = new FileInputStream(onDiskFile);
 243                         try {
 244                                 NB.readFully(in, out, 0, (int) len);
 245                         } finally {
 246                                 in.close();
 247                         }
 248                 }
 249                 return out;
 250         }
 251
 252         /**
 253          * Send this buffer to an output stream.
 254          * <p>
 255          * This method may only be invoked after {@link #close()} has completed
 256          * normally, to ensure all data is completely transferred.
 257          *
 258          * @param os
 259          *            stream to send this buffer's complete content to.
 260          * @param pm
 261          *            if not null progress updates are sent here. Caller should
 262          *            initialize the task and the number of work units to
 263          *            <code>{@link #length()}/1024</code>.
 264          * @throws IOException
 265          *             an error occurred reading from a temporary file on the local
 266          *             system, or writing to the output stream.
 267          */
 268         public void writeTo(final OutputStream os, ProgressMonitor pm)
 269                         throws IOException {
 270                 if (pm == null)
 271                         pm = NullProgressMonitor.INSTANCE;
 272                 if (blocks != null) {
 273                         // Everything is in core so we can stream directly to the output.
 274                         //
 275                         for (final Block b : blocks) {
 276                                 os.write(b.buffer, 0, b.count);
 277                                 pm.update(b.count / 1024);
 278                         }
 279                 } else {
 280                         // Reopen the temporary file and copy the contents.
 281                         //
 282                         final FileInputStream in = new FileInputStream(onDiskFile);
 283                         try {
 284                                 int cnt;
 285                                 final byte[] buf = new byte[Block.SZ];
 286                                 while ((cnt = in.read(buf)) >= 0) {
 287                                         os.write(buf, 0, cnt);
 288                                         pm.update(cnt / 1024);
 289                                 }
 290                         } finally {
 291                                 in.close();
 292                         }
 293                 }
 294         }
 295
 296         /** Clear this buffer so it has no data, and cannot be used again. */
 297         public void destroy() {
 298                 blocks = null;
 299
 300                 if (diskOut != null) {
 301                         try {
 302                                 diskOut.close();
 303                         } catch (IOException err) {
 304                                 // We shouldn't encounter an error closing the file.
 305                         } finally {
 306                                 diskOut = null;
 307                         }
 308                 }
 309
 310                 if (onDiskFile != null) {
 311                         if (!onDiskFile.delete())
 312                                 onDiskFile.deleteOnExit();
 313                         onDiskFile = null;
 314                 }
 315         }
 316
 317         static class Block {
 318                 static final int SZ = 8 * 1024;
 319
 320                 final byte[] buffer = new byte[SZ];
 321
 322                 int count;
 323
 324                 boolean isFull() {
 325                         return count == SZ;
 326                 }
 327         }
 328 }