1/*
2 * Copyright (C) 2008 The Guava Authors
3 *
4 * Licensed under the Apache License, Version 2.0 (the "License");
5 * you may not use this file except in compliance with the License.
6 * You may obtain a copy of the License at
7 *
8 * http://www.apache.org/licenses/LICENSE-2.0
9 *
10 * Unless required by applicable law or agreed to in writing, software
11 * distributed under the License is distributed on an "AS IS" BASIS,
12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13 * See the License for the specific language governing permissions and
14 * limitations under the License.
15 */
16
17package com.google.common.io;
18
19import com.google.common.annotations.Beta;
20import com.google.common.annotations.VisibleForTesting;
21
22import java.io.ByteArrayInputStream;
23import java.io.ByteArrayOutputStream;
24import java.io.File;
25import java.io.FileInputStream;
26import java.io.FileOutputStream;
27import java.io.IOException;
28import java.io.InputStream;
29import java.io.OutputStream;
30
31/**
32 * An {@link OutputStream} that starts buffering to a byte array, but
33 * switches to file buffering once the data reaches a configurable size.
34 *
35 * <p>This class is thread-safe.
36 *
37 * @author Chris Nokleberg
38 * @since 1.0
39 */
40@Beta
41public final class FileBackedOutputStream extends OutputStream {
42
43  private final int fileThreshold;
44  private final boolean resetOnFinalize;
45  private final ByteSource source;
46
47  private OutputStream out;
48  private MemoryOutput memory;
49  private File file;
50
51  /** ByteArrayOutputStream that exposes its internals. */
52  private static class MemoryOutput extends ByteArrayOutputStream {
53    byte[] getBuffer() {
54      return buf;
55    }
56
57    int getCount() {
58      return count;
59    }
60  }
61
62  /** Returns the file holding the data (possibly null). */
63  @VisibleForTesting synchronized File getFile() {
64    return file;
65  }
66
67  /**
68   * Creates a new instance that uses the given file threshold, and does
69   * not reset the data when the {@link ByteSource} returned by
70   * {@link #asByteSource} is finalized.
71   *
72   * @param fileThreshold the number of bytes before the stream should
73   *     switch to buffering to a file
74   */
75  public FileBackedOutputStream(int fileThreshold) {
76    this(fileThreshold, false);
77  }
78
79  /**
80   * Creates a new instance that uses the given file threshold, and
81   * optionally resets the data when the {@link ByteSource} returned
82   * by {@link #asByteSource} is finalized.
83   *
84   * @param fileThreshold the number of bytes before the stream should
85   *     switch to buffering to a file
86   * @param resetOnFinalize if true, the {@link #reset} method will
87   *     be called when the {@link ByteSource} returned by {@link
88   *     #asByteSource} is finalized
89   */
90  public FileBackedOutputStream(int fileThreshold, boolean resetOnFinalize) {
91    this.fileThreshold = fileThreshold;
92    this.resetOnFinalize = resetOnFinalize;
93    memory = new MemoryOutput();
94    out = memory;
95
96    if (resetOnFinalize) {
97      source = new ByteSource() {
98        @Override
99        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 supplier that may be used to retrieve the data buffered
123   * by this stream. This method returns the same object as
124   * {@link #asByteSource()}.
125   *
126   * @deprecated Use {@link #asByteSource()} instead. This method is scheduled
127   *     to be removed in Guava 16.0.
128   */
129  @Deprecated
130  public InputSupplier<InputStream> getSupplier() {
131    return asByteSource();
132  }
133
134  /**
135   * Returns a readable {@link ByteSource} view of the data that has been
136   * written to this stream.
137   *
138   * @since 15.0
139   */
140  public ByteSource asByteSource() {
141    return source;
142  }
143
144  private synchronized InputStream openInputStream() throws IOException {
145    if (file != null) {
146      return new FileInputStream(file);
147    } else {
148      return new ByteArrayInputStream(
149          memory.getBuffer(), 0, memory.getCount());
150    }
151  }
152
153  /**
154   * Calls {@link #close} if not already closed, and then resets this
155   * object back to its initial state, for reuse. If data was buffered
156   * to a file, it will be deleted.
157   *
158   * @throws IOException if an I/O error occurred while deleting the file buffer
159   */
160  public synchronized void reset() throws IOException {
161    try {
162      close();
163    } finally {
164      if (memory == null) {
165        memory = new MemoryOutput();
166      } else {
167        memory.reset();
168      }
169      out = memory;
170      if (file != null) {
171        File deleteMe = file;
172        file = null;
173        if (!deleteMe.delete()) {
174          throw new IOException("Could not delete: " + deleteMe);
175        }
176      }
177    }
178  }
179
180  @Override public synchronized void write(int b) throws IOException {
181    update(1);
182    out.write(b);
183  }
184
185  @Override public synchronized void write(byte[] b) throws IOException {
186    write(b, 0, b.length);
187  }
188
189  @Override public synchronized void write(byte[] b, int off, int len)
190      throws IOException {
191    update(len);
192    out.write(b, off, len);
193  }
194
195  @Override public synchronized void close() throws IOException {
196    out.close();
197  }
198
199  @Override public synchronized void flush() throws IOException {
200    out.flush();
201  }
202
203  /**
204   * Checks if writing {@code len} bytes would go over threshold, and
205   * switches to file buffering if so.
206   */
207  private void update(int len) throws IOException {
208    if (file == null && (memory.getCount() + len > fileThreshold)) {
209      File temp = File.createTempFile("FileBackedOutputStream", null);
210      if (resetOnFinalize) {
211        // Finalizers are not guaranteed to be called on system shutdown;
212        // this is insurance.
213        temp.deleteOnExit();
214      }
215      FileOutputStream transfer = new FileOutputStream(temp);
216      transfer.write(memory.getBuffer(), 0, memory.getCount());
217      transfer.flush();
218
219      // We've successfully transferred the data; switch to writing to file
220      out = transfer;
221      file = temp;
222      memory = null;
223    }
224  }
225}
226