001/*
002 * Copyright (C) 2008 The Guava 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 * 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
017package com.google.common.io;
018
019import com.google.common.annotations.Beta;
020import com.google.common.annotations.VisibleForTesting;
021
022import java.io.ByteArrayInputStream;
023import java.io.ByteArrayOutputStream;
024import java.io.File;
025import java.io.FileInputStream;
026import java.io.FileOutputStream;
027import java.io.IOException;
028import java.io.InputStream;
029import java.io.OutputStream;
030
031/**
032 * An {@link OutputStream} that starts buffering to a byte array, but
033 * switches to file buffering once the data reaches a configurable size.
034 *
035 * <p>This class is thread-safe.
036 *
037 * @author Chris Nokleberg
038 * @since 1
039 */
040@Beta
041public final class FileBackedOutputStream extends OutputStream {
042
043  private final int fileThreshold;
044  private final boolean resetOnFinalize;
045  private final InputSupplier<InputStream> supplier;
046
047  private OutputStream out;
048  private MemoryOutput memory;
049  private File file;
050
051  /** ByteArrayOutputStream that exposes its internals. */
052  private static class MemoryOutput extends ByteArrayOutputStream {
053    byte[] getBuffer() {
054      return buf;
055    }
056
057    int getCount() {
058      return count;
059    }
060  }
061
062  /** Returns the file holding the data (possibly null). */
063  @VisibleForTesting synchronized File getFile() {
064    return file;
065  }
066
067  /**
068   * Creates a new instance that uses the given file threshold.
069   * Equivalent to {@code ThresholdOutputStream(fileThreshold, false)}.
070   *
071   * @param fileThreshold the number of bytes before the stream should
072   *     switch to buffering to a file
073   */
074  public FileBackedOutputStream(int fileThreshold) {
075    this(fileThreshold, false);
076  }
077
078  /**
079   * Creates a new instance that uses the given file threshold, and
080   * optionally resets the data when the {@link InputSupplier} returned
081   * by {@link #getSupplier} is finalized.
082   *
083   * @param fileThreshold the number of bytes before the stream should
084   *     switch to buffering to a file
085   * @param resetOnFinalize if true, the {@link #reset} method will
086   *     be called when the {@link InputSupplier} returned by {@link
087   *     #getSupplier} is finalized
088   */
089  public FileBackedOutputStream(int fileThreshold, boolean resetOnFinalize) {
090    this.fileThreshold = fileThreshold;
091    this.resetOnFinalize = resetOnFinalize;
092    memory = new MemoryOutput();
093    out = memory;
094
095    if (resetOnFinalize) {
096      supplier = new InputSupplier<InputStream>() {
097        @Override
098        public InputStream getInput() throws IOException {
099          return openStream();
100        }
101
102        @Override protected void finalize() {
103          try {
104            reset();
105          } catch (Throwable t) {
106            t.printStackTrace(System.err);
107          }
108        }
109      };
110    } else {
111      supplier = new InputSupplier<InputStream>() {
112        @Override
113        public InputStream getInput() throws IOException {
114          return openStream();
115        }
116      };
117    }
118  }
119
120  /**
121   * Returns a supplier that may be used to retrieve the data buffered
122   * by this stream.
123   */
124  public InputSupplier<InputStream> getSupplier() {
125    return supplier;
126  }
127
128  private synchronized InputStream openStream() throws IOException {
129    if (file != null) {
130      return new FileInputStream(file);
131    } else {
132      return new ByteArrayInputStream(
133          memory.getBuffer(), 0, memory.getCount());
134    }
135  }
136
137  /**
138   * Calls {@link #close} if not already closed, and then resets this
139   * object back to its initial state, for reuse. If data was buffered
140   * to a file, it will be deleted.
141   *
142   * @throws IOException if an I/O error occurred while deleting the file buffer
143   */
144  public synchronized void reset() throws IOException {
145    try {
146      close();
147    } finally {
148      if (memory == null) {
149        memory = new MemoryOutput();
150      } else {
151        memory.reset();
152      }
153      out = memory;
154      if (file != null) {
155        File deleteMe = file;
156        file = null;
157        if (!deleteMe.delete()) {
158          throw new IOException("Could not delete: " + deleteMe);
159        }
160      }
161    }
162  }
163
164  @Override public synchronized void write(int b) throws IOException {
165    update(1);
166    out.write(b);
167  }
168
169  @Override public synchronized void write(byte[] b) throws IOException {
170    write(b, 0, b.length);
171  }
172
173  @Override public synchronized void write(byte[] b, int off, int len)
174      throws IOException {
175    update(len);
176    out.write(b, off, len);
177  }
178
179  @Override public synchronized void close() throws IOException {
180    out.close();
181  }
182
183  @Override public synchronized void flush() throws IOException {
184    out.flush();
185  }
186
187  /**
188   * Checks if writing {@code len} bytes would go over threshold, and
189   * switches to file buffering if so.
190   */
191  private void update(int len) throws IOException {
192    if (file == null && (memory.getCount() + len > fileThreshold)) {
193      File temp = File.createTempFile("FileBackedOutputStream", null);
194      if (resetOnFinalize) {
195        // Finalizers are not guaranteed to be called on system shutdown;
196        // this is insurance.
197        temp.deleteOnExit();
198      }
199      FileOutputStream transfer = new FileOutputStream(temp);
200      transfer.write(memory.getBuffer(), 0, memory.getCount());
201      transfer.flush();
202
203      // We've successfully transferred the data; switch to writing to file
204      out = transfer;
205      file = temp;
206      memory = null;
207    }
208  }
209}