001/*
002 * Copyright (C) 2008 The Guava Authors
003 *
004 * Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
005 * in compliance with the License. You may obtain a copy of the License at
006 *
007 * http://www.apache.org/licenses/LICENSE-2.0
008 *
009 * Unless required by applicable law or agreed to in writing, software distributed under the License
010 * is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
011 * or implied. See the License for the specific language governing permissions and limitations under
012 * the License.
013 */
014
015package com.google.common.io;
016
017import com.google.common.annotations.Beta;
018import com.google.common.annotations.GwtIncompatible;
019import com.google.common.annotations.VisibleForTesting;
020import java.io.ByteArrayInputStream;
021import java.io.ByteArrayOutputStream;
022import java.io.File;
023import java.io.FileInputStream;
024import java.io.FileOutputStream;
025import java.io.IOException;
026import java.io.InputStream;
027import java.io.OutputStream;
028import org.checkerframework.checker.nullness.qual.Nullable;
029
030/**
031 * An {@link OutputStream} that starts buffering to a byte array, but switches to file buffering
032 * once the data reaches a configurable size.
033 *
034 * <p>This class is thread-safe.
035 *
036 * @author Chris Nokleberg
037 * @since 1.0
038 */
039@Beta
040@GwtIncompatible
041public final class FileBackedOutputStream extends OutputStream {
042
043  private final int fileThreshold;
044  private final boolean resetOnFinalize;
045  private final ByteSource source;
046
047  private OutputStream out;
048  private MemoryOutput memory;
049  private @Nullable 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
064  synchronized File getFile() {
065    return file;
066  }
067
068  /**
069   * Creates a new instance that uses the given file threshold, and does not reset the data when the
070   * {@link ByteSource} returned by {@link #asByteSource} is finalized.
071   *
072   * @param fileThreshold the number of bytes before the stream should 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 optionally resets the data when
080   * the {@link ByteSource} returned by {@link #asByteSource} is finalized.
081   *
082   * @param fileThreshold the number of bytes before the stream should switch to buffering to a file
083   * @param resetOnFinalize if true, the {@link #reset} method will be called when the {@link
084   *     ByteSource} returned by {@link #asByteSource} is finalized
085   */
086  public FileBackedOutputStream(int fileThreshold, boolean resetOnFinalize) {
087    this.fileThreshold = fileThreshold;
088    this.resetOnFinalize = resetOnFinalize;
089    memory = new MemoryOutput();
090    out = memory;
091
092    if (resetOnFinalize) {
093      source =
094          new ByteSource() {
095            @Override
096            public InputStream openStream() throws IOException {
097              return openInputStream();
098            }
099
100            @Override
101            protected void finalize() {
102              try {
103                reset();
104              } catch (Throwable t) {
105                t.printStackTrace(System.err);
106              }
107            }
108          };
109    } else {
110      source =
111          new ByteSource() {
112            @Override
113            public InputStream openStream() throws IOException {
114              return openInputStream();
115            }
116          };
117    }
118  }
119
120  /**
121   * Returns a readable {@link ByteSource} view of the data that has been written to this stream.
122   *
123   * @since 15.0
124   */
125  public ByteSource asByteSource() {
126    return source;
127  }
128
129  private synchronized InputStream openInputStream() throws IOException {
130    if (file != null) {
131      return new FileInputStream(file);
132    } else {
133      return new ByteArrayInputStream(memory.getBuffer(), 0, memory.getCount());
134    }
135  }
136
137  /**
138   * Calls {@link #close} if not already closed, and then resets this object back to its initial
139   * state, for reuse. If data was buffered to a file, it will be deleted.
140   *
141   * @throws IOException if an I/O error occurred while deleting the file buffer
142   */
143  public synchronized void reset() throws IOException {
144    try {
145      close();
146    } finally {
147      if (memory == null) {
148        memory = new MemoryOutput();
149      } else {
150        memory.reset();
151      }
152      out = memory;
153      if (file != null) {
154        File deleteMe = file;
155        file = null;
156        if (!deleteMe.delete()) {
157          throw new IOException("Could not delete: " + deleteMe);
158        }
159      }
160    }
161  }
162
163  @Override
164  public synchronized void write(int b) throws IOException {
165    update(1);
166    out.write(b);
167  }
168
169  @Override
170  public synchronized void write(byte[] b) throws IOException {
171    write(b, 0, b.length);
172  }
173
174  @Override
175  public synchronized void write(byte[] b, int off, int len) throws IOException {
176    update(len);
177    out.write(b, off, len);
178  }
179
180  @Override
181  public synchronized void close() throws IOException {
182    out.close();
183  }
184
185  @Override
186  public synchronized void flush() throws IOException {
187    out.flush();
188  }
189
190  /**
191   * Checks if writing {@code len} bytes would go over threshold, and switches to file buffering if
192   * so.
193   */
194  private void update(int len) throws IOException {
195    if (file == null && (memory.getCount() + len > fileThreshold)) {
196      File temp = File.createTempFile("FileBackedOutputStream", null);
197      if (resetOnFinalize) {
198        // Finalizers are not guaranteed to be called on system shutdown;
199        // this is insurance.
200        temp.deleteOnExit();
201      }
202      FileOutputStream transfer = new FileOutputStream(temp);
203      transfer.write(memory.getBuffer(), 0, memory.getCount());
204      transfer.flush();
205
206      // We've successfully transferred the data; switch to writing to file
207      out = transfer;
208      file = temp;
209      memory = null;
210    }
211  }
212}