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
017
018import com.google.common.annotations.Beta;
019import com.google.common.annotations.GwtIncompatible;
020import com.google.common.annotations.VisibleForTesting;
021import com.google.errorprone.annotations.concurrent.GuardedBy;
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;
030import org.checkerframework.checker.nullness.qual.Nullable;
031
032/**
033 * An {@link OutputStream} that starts buffering to a byte array, but switches to file buffering
034 * once the data reaches a configurable size.
035 *
036 * <p>Temporary files created by this stream may live in the local filesystem until either:
037 *
038 * <ul>
039 *   <li>{@link #reset} is called (removing the data in this stream and deleting the file), or...
040 *   <li>this stream (or, more precisely, its {@link #asByteSource} view) is finalized during
041 *       garbage collection, <strong>AND</strong> this stream was not constructed with {@linkplain
042 *       #FileBackedOutputStream(int) the 1-arg constructor} or the {@linkplain
043 *       #FileBackedOutputStream(int, boolean) 2-arg constructor} passing {@code false} in the
044 *       second parameter.
045 * </ul>
046 *
047 * <p>This class is thread-safe.
048 *
049 * @author Chris Nokleberg
050 * @since 1.0
051 */
052@Beta
053@GwtIncompatible
054public final class FileBackedOutputStream extends OutputStream {
055  private final int fileThreshold;
056  private final boolean resetOnFinalize;
057  private final ByteSource source;
058  @Nullable private final File parentDirectory;
059
060  @GuardedBy("this")
061  private OutputStream out;
062
063  @GuardedBy("this")
064  private MemoryOutput memory;
065
066  @GuardedBy("this")
067  private @Nullable File file;
068
069  /** ByteArrayOutputStream that exposes its internals. */
070  private static class MemoryOutput extends ByteArrayOutputStream {
071    byte[] getBuffer() {
072      return buf;
073    }
074
075    int getCount() {
076      return count;
077    }
078  }
079
080  /** Returns the file holding the data (possibly null). */
081  @VisibleForTesting
082  synchronized File getFile() {
083    return file;
084  }
085
086  /**
087   * Creates a new instance that uses the given file threshold, and does not reset the data when the
088   * {@link ByteSource} returned by {@link #asByteSource} is finalized.
089   *
090   * @param fileThreshold the number of bytes before the stream should switch to buffering to a file
091   */
092  public FileBackedOutputStream(int fileThreshold) {
093    this(fileThreshold, false);
094  }
095
096  /**
097   * Creates a new instance that uses the given file threshold, and optionally resets the data when
098   * the {@link ByteSource} returned by {@link #asByteSource} is finalized.
099   *
100   * @param fileThreshold the number of bytes before the stream should switch to buffering to a file
101   * @param resetOnFinalize if true, the {@link #reset} method will be called when the {@link
102   *     ByteSource} returned by {@link #asByteSource} is finalized.
103   */
104  public FileBackedOutputStream(int fileThreshold, boolean resetOnFinalize) {
105    this(fileThreshold, resetOnFinalize, null);
106  }
107
108  private FileBackedOutputStream(
109      int fileThreshold, boolean resetOnFinalize, @Nullable File parentDirectory) {
110    this.fileThreshold = fileThreshold;
111    this.resetOnFinalize = resetOnFinalize;
112    this.parentDirectory = parentDirectory;
113    memory = new MemoryOutput();
114    out = memory;
115
116    if (resetOnFinalize) {
117      source =
118          new ByteSource() {
119            @Override
120            public InputStream openStream() throws IOException {
121              return openInputStream();
122            }
123
124            @Override
125            protected void finalize() {
126              try {
127                reset();
128              } catch (Throwable t) {
129                t.printStackTrace(System.err);
130              }
131            }
132          };
133    } else {
134      source =
135          new ByteSource() {
136            @Override
137            public InputStream openStream() throws IOException {
138              return openInputStream();
139            }
140          };
141    }
142  }
143
144  /**
145   * Returns a readable {@link ByteSource} view of the data that has been written to this stream.
146   *
147   * @since 15.0
148   */
149  public ByteSource asByteSource() {
150    return source;
151  }
152
153  private synchronized InputStream openInputStream() throws IOException {
154    if (file != null) {
155      return new FileInputStream(file);
156    } else {
157      return new ByteArrayInputStream(memory.getBuffer(), 0, memory.getCount());
158    }
159  }
160
161  /**
162   * Calls {@link #close} if not already closed, and then resets this object back to its initial
163   * state, for reuse. If data was buffered to a file, it will be deleted.
164   *
165   * @throws IOException if an I/O error occurred while deleting the file buffer
166   */
167  public synchronized void reset() throws IOException {
168    try {
169      close();
170    } finally {
171      if (memory == null) {
172        memory = new MemoryOutput();
173      } else {
174        memory.reset();
175      }
176      out = memory;
177      if (file != null) {
178        File deleteMe = file;
179        file = null;
180        if (!deleteMe.delete()) {
181          throw new IOException("Could not delete: " + deleteMe);
182        }
183      }
184    }
185  }
186
187  @Override
188  public synchronized void write(int b) throws IOException {
189    update(1);
190    out.write(b);
191  }
192
193  @Override
194  public synchronized void write(byte[] b) throws IOException {
195    write(b, 0, b.length);
196  }
197
198  @Override
199  public synchronized void write(byte[] b, int off, int len) throws IOException {
200    update(len);
201    out.write(b, off, len);
202  }
203
204  @Override
205  public synchronized void close() throws IOException {
206    out.close();
207  }
208
209  @Override
210  public synchronized void flush() throws IOException {
211    out.flush();
212  }
213
214  /**
215   * Checks if writing {@code len} bytes would go over threshold, and switches to file buffering if
216   * so.
217   */
218  @GuardedBy("this")
219  private void update(int len) throws IOException {
220    if (file == null && (memory.getCount() + len > fileThreshold)) {
221      File temp = File.createTempFile("FileBackedOutputStream", null, parentDirectory);
222      if (resetOnFinalize) {
223        // Finalizers are not guaranteed to be called on system shutdown;
224        // this is insurance.
225        temp.deleteOnExit();
226      }
227      FileOutputStream transfer = new FileOutputStream(temp);
228      transfer.write(memory.getBuffer(), 0, memory.getCount());
229      transfer.flush();
230
231      // We've successfully transferred the data; switch to writing to file
232      out = transfer;
233      file = temp;
234      memory = null;
235    }
236  }
237}