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.0
039 */
040@Beta
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 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, and does
069   * not reset the data when the {@link ByteSource} returned by
070   * {@link #asByteSource} is finalized.
071   *
072   * @param fileThreshold the number of bytes before the stream should
073   *     switch to buffering to a file
074   */
075  public FileBackedOutputStream(int fileThreshold) {
076    this(fileThreshold, false);
077  }
078
079  /**
080   * Creates a new instance that uses the given file threshold, and
081   * optionally resets the data when the {@link ByteSource} returned
082   * by {@link #asByteSource} is finalized.
083   *
084   * @param fileThreshold the number of bytes before the stream should
085   *     switch to buffering to a file
086   * @param resetOnFinalize if true, the {@link #reset} method will
087   *     be called when the {@link ByteSource} returned by {@link
088   *     #asByteSource} is finalized
089   */
090  public FileBackedOutputStream(int fileThreshold, boolean resetOnFinalize) {
091    this.fileThreshold = fileThreshold;
092    this.resetOnFinalize = resetOnFinalize;
093    memory = new MemoryOutput();
094    out = memory;
095
096    if (resetOnFinalize) {
097      source = new ByteSource() {
098        @Override
099        public InputStream openStream() throws IOException {
100          return openInputStream();
101        }
102
103        @Override protected void finalize() {
104          try {
105            reset();
106          } catch (Throwable t) {
107            t.printStackTrace(System.err);
108          }
109        }
110      };
111    } else {
112      source = new ByteSource() {
113        @Override
114        public InputStream openStream() throws IOException {
115          return openInputStream();
116        }
117      };
118    }
119  }
120
121  /**
122   * Returns a readable {@link ByteSource} view of the data that has been
123   * written to this stream.
124   *
125   * @since 15.0
126   */
127  public ByteSource asByteSource() {
128    return source;
129  }
130
131  private synchronized InputStream openInputStream() throws IOException {
132    if (file != null) {
133      return new FileInputStream(file);
134    } else {
135      return new ByteArrayInputStream(
136          memory.getBuffer(), 0, memory.getCount());
137    }
138  }
139
140  /**
141   * Calls {@link #close} if not already closed, and then resets this
142   * object back to its initial state, for reuse. If data was buffered
143   * to a file, it will be deleted.
144   *
145   * @throws IOException if an I/O error occurred while deleting the file buffer
146   */
147  public synchronized void reset() throws IOException {
148    try {
149      close();
150    } finally {
151      if (memory == null) {
152        memory = new MemoryOutput();
153      } else {
154        memory.reset();
155      }
156      out = memory;
157      if (file != null) {
158        File deleteMe = file;
159        file = null;
160        if (!deleteMe.delete()) {
161          throw new IOException("Could not delete: " + deleteMe);
162        }
163      }
164    }
165  }
166
167  @Override public synchronized void write(int b) throws IOException {
168    update(1);
169    out.write(b);
170  }
171
172  @Override public synchronized void write(byte[] b) throws IOException {
173    write(b, 0, b.length);
174  }
175
176  @Override public synchronized void write(byte[] b, int off, int len)
177      throws IOException {
178    update(len);
179    out.write(b, off, len);
180  }
181
182  @Override public synchronized void close() throws IOException {
183    out.close();
184  }
185
186  @Override public synchronized void flush() throws IOException {
187    out.flush();
188  }
189
190  /**
191   * Checks if writing {@code len} bytes would go over threshold, and
192   * switches to file buffering if 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}