    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
    017    package com.google.common.io;
    018
    019    import com.google.common.annotations.Beta;
    020    import com.google.common.annotations.VisibleForTesting;
    021
    022    import java.io.ByteArrayInputStream;
    023    import java.io.ByteArrayOutputStream;
    024    import java.io.File;
    025    import java.io.FileInputStream;
    026    import java.io.FileOutputStream;
    027    import java.io.IOException;
    028    import java.io.InputStream;
    029    import 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
    041    public final class FileBackedOutputStream extends OutputStream {
    042
    043      private final int fileThreshold;
    044      private final boolean resetOnFinalize;
    045      private final InputSupplier<InputStream> supplier;
    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 InputSupplier} returned by
    070       * {@link #getSupplier} 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 InputSupplier} returned
    082       * by {@link #getSupplier} 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 InputSupplier} returned by {@link
    088       *     #getSupplier} 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          supplier = new InputSupplier<InputStream>() {
    098            @Override
    099            public InputStream getInput() throws IOException {
    100              return openStream();
    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          supplier = new InputSupplier<InputStream>() {
    113            @Override
    114            public InputStream getInput() throws IOException {
    115              return openStream();
    116            }
    117          };
    118        }
    119      }
    120
    121      /**
    122       * Returns a supplier that may be used to retrieve the data buffered
    123       * by this stream.
    124       */
    125      public InputSupplier<InputStream> getSupplier() {
    126        return supplier;
    127      }
    128
    129      private synchronized InputStream openStream() throws IOException {
    130        if (file != null) {
    131          return new FileInputStream(file);
    132        } else {
    133          return new ByteArrayInputStream(
    134              memory.getBuffer(), 0, memory.getCount());
    135        }
    136      }
    137
    138      /**
    139       * Calls {@link #close} if not already closed, and then resets this
    140       * object back to its initial state, for reuse. If data was buffered
    141       * to a file, it will be deleted.
    142       *
    143       * @throws IOException if an I/O error occurred while deleting the file buffer
    144       */
    145      public synchronized void reset() throws IOException {
    146        try {
    147          close();
    148        } finally {
    149          if (memory == null) {
    150            memory = new MemoryOutput();
    151          } else {
    152            memory.reset();
    153          }
    154          out = memory;
    155          if (file != null) {
    156            File deleteMe = file;
    157            file = null;
    158            if (!deleteMe.delete()) {
    159              throw new IOException("Could not delete: " + deleteMe);
    160            }
    161          }
    162        }
    163      }
    164
    165      @Override public synchronized void write(int b) throws IOException {
    166        update(1);
    167        out.write(b);
    168      }
    169
    170      @Override public synchronized void write(byte[] b) throws IOException {
    171        write(b, 0, b.length);
    172      }
    173
    174      @Override public synchronized void write(byte[] b, int off, int len)
    175          throws IOException {
    176        update(len);
    177        out.write(b, off, len);
    178      }
    179
    180      @Override public synchronized void close() throws IOException {
    181        out.close();
    182      }
    183
    184      @Override public synchronized void flush() throws IOException {
    185        out.flush();
    186      }
    187
    188      /**
    189       * Checks if writing {@code len} bytes would go over threshold, and
    190       * switches to file buffering if so.
    191       */
    192      private void update(int len) throws IOException {
    193        if (file == null && (memory.getCount() + len > fileThreshold)) {
    194          File temp = File.createTempFile("FileBackedOutputStream", null);
    195          if (resetOnFinalize) {
    196            // Finalizers are not guaranteed to be called on system shutdown;
    197            // this is insurance.
    198            temp.deleteOnExit();
    199          }
    200          FileOutputStream transfer = new FileOutputStream(temp);
    201          transfer.write(memory.getBuffer(), 0, memory.getCount());
    202          transfer.flush();
    203
    204          // We've successfully transferred the data; switch to writing to file
    205          out = transfer;
    206          file = temp;
    207          memory = null;
    208        }
    209      }
    210    }