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.compatqual.NullableDecl;
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  @NullableDecl 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  @NullableDecl
068  private File file;
069
070  /** ByteArrayOutputStream that exposes its internals. */
071  private static class MemoryOutput extends ByteArrayOutputStream {
072    byte[] getBuffer() {
073      return buf;
074    }
075
076    int getCount() {
077      return count;
078    }
079  }
080
081  /** Returns the file holding the data (possibly null). */
082  @VisibleForTesting
083  synchronized File getFile() {
084    return file;
085  }
086
087  /**
088   * Creates a new instance that uses the given file threshold, and does not reset the data when the
089   * {@link ByteSource} returned by {@link #asByteSource} is finalized.
090   *
091   * @param fileThreshold the number of bytes before the stream should switch to buffering to a file
092   */
093  public FileBackedOutputStream(int fileThreshold) {
094    this(fileThreshold, false);
095  }
096
097  /**
098   * Creates a new instance that uses the given file threshold, and optionally resets the data when
099   * the {@link ByteSource} returned by {@link #asByteSource} is finalized.
100   *
101   * @param fileThreshold the number of bytes before the stream should switch to buffering to a file
102   * @param resetOnFinalize if true, the {@link #reset} method will be called when the {@link
103   *     ByteSource} returned by {@link #asByteSource} is finalized.
104   */
105  public FileBackedOutputStream(int fileThreshold, boolean resetOnFinalize) {
106    this(fileThreshold, resetOnFinalize, null);
107  }
108
109  private FileBackedOutputStream(
110      int fileThreshold, boolean resetOnFinalize, @NullableDecl File parentDirectory) {
111    this.fileThreshold = fileThreshold;
112    this.resetOnFinalize = resetOnFinalize;
113    this.parentDirectory = parentDirectory;
114    memory = new MemoryOutput();
115    out = memory;
116
117    if (resetOnFinalize) {
118      source =
119          new ByteSource() {
120            @Override
121            public InputStream openStream() throws IOException {
122              return openInputStream();
123            }
124
125            @Override
126            protected void finalize() {
127              try {
128                reset();
129              } catch (Throwable t) {
130                t.printStackTrace(System.err);
131              }
132            }
133          };
134    } else {
135      source =
136          new ByteSource() {
137            @Override
138            public InputStream openStream() throws IOException {
139              return openInputStream();
140            }
141          };
142    }
143  }
144
145  /**
146   * Returns a readable {@link ByteSource} view of the data that has been written to this stream.
147   *
148   * @since 15.0
149   */
150  public ByteSource asByteSource() {
151    return source;
152  }
153
154  private synchronized InputStream openInputStream() throws IOException {
155    if (file != null) {
156      return new FileInputStream(file);
157    } else {
158      return new ByteArrayInputStream(memory.getBuffer(), 0, memory.getCount());
159    }
160  }
161
162  /**
163   * Calls {@link #close} if not already closed, and then resets this object back to its initial
164   * state, for reuse. If data was buffered to a file, it will be deleted.
165   *
166   * @throws IOException if an I/O error occurred while deleting the file buffer
167   */
168  public synchronized void reset() throws IOException {
169    try {
170      close();
171    } finally {
172      if (memory == null) {
173        memory = new MemoryOutput();
174      } else {
175        memory.reset();
176      }
177      out = memory;
178      if (file != null) {
179        File deleteMe = file;
180        file = null;
181        if (!deleteMe.delete()) {
182          throw new IOException("Could not delete: " + deleteMe);
183        }
184      }
185    }
186  }
187
188  @Override
189  public synchronized void write(int b) throws IOException {
190    update(1);
191    out.write(b);
192  }
193
194  @Override
195  public synchronized void write(byte[] b) throws IOException {
196    write(b, 0, b.length);
197  }
198
199  @Override
200  public synchronized void write(byte[] b, int off, int len) throws IOException {
201    update(len);
202    out.write(b, off, len);
203  }
204
205  @Override
206  public synchronized void close() throws IOException {
207    out.close();
208  }
209
210  @Override
211  public synchronized void flush() throws IOException {
212    out.flush();
213  }
214
215  /**
216   * Checks if writing {@code len} bytes would go over threshold, and switches to file buffering if
217   * so.
218   */
219  @GuardedBy("this")
220  private void update(int len) throws IOException {
221    if (file == null && (memory.getCount() + len > fileThreshold)) {
222      File temp = File.createTempFile("FileBackedOutputStream", null, parentDirectory);
223      if (resetOnFinalize) {
224        // Finalizers are not guaranteed to be called on system shutdown;
225        // this is insurance.
226        temp.deleteOnExit();
227      }
228      FileOutputStream transfer = new FileOutputStream(temp);
229      transfer.write(memory.getBuffer(), 0, memory.getCount());
230      transfer.flush();
231
232      // We've successfully transferred the data; switch to writing to file
233      out = transfer;
234      file = temp;
235      memory = null;
236    }
237  }
238}