    001/*
    002 * Licensed to the Apache Software Foundation (ASF) under one or more
    003 * contributor license agreements.  See the NOTICE file distributed with
    004 * this work for additional information regarding copyright ownership.
    005 * The ASF licenses this file to You under the Apache License, Version 2.0
    006 * (the "License"); you may not use this file except in compliance with
    007 * the License.  You may obtain a copy of the License at
    008 *
    009 *      https://www.apache.org/licenses/LICENSE-2.0
    010 *
    011 * Unless required by applicable law or agreed to in writing, software
    012 * distributed under the License is distributed on an "AS IS" BASIS,
    013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    014 * See the License for the specific language governing permissions and
    015 * limitations under the License.
    016 */
    017package org.apache.commons.io.output;
    018
    019import java.io.IOException;
    020import java.io.OutputStream;
    021
    022import org.apache.commons.io.IOUtils;
    023import org.apache.commons.io.function.IOConsumer;
    024import org.apache.commons.io.function.IOFunction;
    025
    026/**
    027 * An output stream which triggers an event on the first write that causes
    028 * the total number of bytes written to the stream to exceed a configured threshold,
    029 * and every subsequent write. The event
    030 * can be used, for example, to throw an exception if a maximum has been reached,
    031 * or to switch the underlying stream when the threshold is exceeded.
    032 *
    033 * <p>
    034 * This class overrides all {@link OutputStream} methods. However, these overrides ultimately call the corresponding
    035 * methods in the underlying output stream implementation.
    036 * </p>
    037 * <p>
    038 * NOTE: This implementation may trigger the event <em>before</em> the threshold is actually reached, since it triggers
    039 * when a pending write operation would cause the threshold to be exceeded.
    040 * </p>
    041 * <p>
    042 * See also the subclass {@link DeferredFileOutputStream}.
    043 * </p>
    044 *
    045 * @see DeferredFileOutputStream
    046 */
    047public class ThresholdingOutputStream extends OutputStream {
    048
    049    /**
    050     * Noop output stream getter function.
    051     */
    052    private static final IOFunction<ThresholdingOutputStream, OutputStream> NOOP_OS_GETTER = os -> NullOutputStream.INSTANCE;
    053
    054    /**
    055     * The threshold at which the event will be triggered.
    056     */
    057    private final int threshold;
    058
    059    /**
    060     * Accepts reaching the threshold.
    061     */
    062    private final IOConsumer<ThresholdingOutputStream> thresholdConsumer;
    063
    064    /**
    065     * Gets the output stream.
    066     */
    067    private final IOFunction<ThresholdingOutputStream, OutputStream> outputStreamGetter;
    068
    069    /**
    070     * The number of bytes written to the output stream.
    071     */
    072    private long written;
    073
    074    /**
    075     * Whether or not the configured threshold has been exceeded.
    076     */
    077    private boolean thresholdExceeded;
    078
    079    /**
    080     * Constructs an instance of this class which will trigger an event at the specified threshold.
    081     *
    082     * @param threshold The number of bytes at which to trigger an event.
    083     */
    084    public ThresholdingOutputStream(final int threshold) {
    085        this(threshold, IOConsumer.noop(), NOOP_OS_GETTER);
    086    }
    087
    088    /**
    089     * Constructs an instance of this class which will trigger an event at the specified threshold.
    090     * A negative threshold has no meaning and will be treated as 0
    091     *
    092     * @param threshold The number of bytes at which to trigger an event.
    093     * @param thresholdConsumer Accepts reaching the threshold.
    094     * @param outputStreamGetter Gets the output stream.
    095     * @since 2.9.0
    096     */
    097    public ThresholdingOutputStream(final int threshold, final IOConsumer<ThresholdingOutputStream> thresholdConsumer,
    098        final IOFunction<ThresholdingOutputStream, OutputStream> outputStreamGetter) {
    099        this.threshold = threshold < 0 ? 0 : threshold;
    100        this.thresholdConsumer = thresholdConsumer == null ? IOConsumer.noop() : thresholdConsumer;
    101        this.outputStreamGetter = outputStreamGetter == null ? NOOP_OS_GETTER : outputStreamGetter;
    102    }
    103
    104    /**
    105     * Checks to see if writing the specified number of bytes would cause the configured threshold to be exceeded. If
    106     * so, triggers an event to allow a concrete implementation to take action on this.
    107     *
    108     * @param count The number of bytes about to be written to the underlying output stream.
    109     * @throws IOException if an error occurs.
    110     */
    111    protected void checkThreshold(final int count) throws IOException {
    112        if (!thresholdExceeded && written + count > threshold) {
    113            thresholdExceeded = true;
    114            thresholdReached();
    115        }
    116    }
    117
    118    /**
    119     * Closes this output stream and releases any system resources associated with this stream.
    120     *
    121     * @throws IOException if an error occurs.
    122     */
    123    @Override
    124    public void close() throws IOException {
    125        try {
    126            flush();
    127        } catch (final IOException ignored) {
    128            // ignore
    129        }
    130        // TODO for 4.0: Replace with getOutputStream()
    131        getStream().close();
    132    }
    133
    134    /**
    135     * Flushes this output stream and forces any buffered output bytes to be written out.
    136     *
    137     * @throws IOException if an error occurs.
    138     */
    139    @SuppressWarnings("resource") // the underlying stream is managed by a subclass.
    140    @Override
    141    public void flush() throws IOException {
    142        // TODO for 4.0: Replace with getOutputStream()
    143        getStream().flush();
    144    }
    145
    146    /**
    147     * Gets the number of bytes that have been written to this output stream.
    148     *
    149     * @return The number of bytes written.
    150     */
    151    public long getByteCount() {
    152        return written;
    153    }
    154
    155    /**
    156     * Gets the underlying output stream, to which the corresponding {@link OutputStream} methods in this class will
    157     * ultimately delegate.
    158     *
    159     * @return The underlying output stream.
    160     * @throws IOException if an error occurs.
    161     * @since 2.14.0
    162     */
    163    protected OutputStream getOutputStream() throws IOException {
    164        return outputStreamGetter.apply(this);
    165    }
    166
    167    /**
    168     * Gets the underlying output stream, to which the corresponding {@link OutputStream} methods in this class will
    169     * ultimately delegate.
    170     *
    171     * @return The underlying output stream.
    172     * @throws IOException if an error occurs.
    173     * @deprecated Use {@link #getOutputStream()}.
    174     */
    175    @Deprecated
    176    protected OutputStream getStream() throws IOException {
    177        return getOutputStream();
    178    }
    179
    180    /**
    181     * Gets the threshold, in bytes, at which an event will be triggered.
    182     *
    183     * @return The threshold point, in bytes.
    184     */
    185    public int getThreshold() {
    186        return threshold;
    187    }
    188
    189    /**
    190     * Tests whether or not the configured threshold has been exceeded for this output stream.
    191     *
    192     * @return {@code true} if the threshold has been reached; {@code false} otherwise.
    193     */
    194    public boolean isThresholdExceeded() {
    195        return written > threshold;
    196    }
    197
    198    /**
    199     * Resets the byteCount to zero. You can call this from {@link #thresholdReached()} if you want the event to be
    200     * triggered again.
    201     */
    202    protected void resetByteCount() {
    203        this.thresholdExceeded = false;
    204        this.written = 0;
    205    }
    206
    207    /**
    208     * Sets the byteCount to count. Useful for re-opening an output stream that has previously been written to.
    209     *
    210     * @param count The number of bytes that have already been written to the output stream.
    211     * @since 2.5
    212     */
    213    protected void setByteCount(final long count) {
    214        this.written = count;
    215    }
    216
    217    /**
    218     * Indicates that the configured threshold has been reached, and that a subclass should take whatever action
    219     * necessary on this event. This may include changing the underlying output stream.
    220     *
    221     * @throws IOException if an error occurs.
    222     */
    223    protected void thresholdReached() throws IOException {
    224        thresholdConsumer.accept(this);
    225    }
    226
    227    /**
    228     * Writes {@code b.length} bytes from the specified byte array to this output stream.
    229     *
    230     * @param b The array of bytes to be written.
    231     * @throws NullPointerException if the byte array is {@code null}.
    232     * @throws IOException if an error occurs.
    233     */
    234    @SuppressWarnings("resource") // the underlying stream is managed by a subclass.
    235    @Override
    236    public void write(final byte[] b) throws IOException {
    237        checkThreshold(b.length);
    238        // TODO for 4.0: Replace with getOutputStream()
    239        getStream().write(b);
    240        written += b.length;
    241    }
    242
    243    /**
    244     * Writes {@code len} bytes from the specified byte array starting at offset {@code off} to this output stream.
    245     *
    246     * @param b The byte array from which the data will be written.
    247     * @param off The start offset in the byte array.
    248     * @param len The number of bytes to write.
    249     * @throws NullPointerException if the byte array is {@code null}.
    250     * @throws IndexOutOfBoundsException if {@code off} or {@code len} are negative, or if {@code off + len} is greater than {@code b.length}.
    251     * @throws IOException if an error occurs.
    252     */
    253    @SuppressWarnings("resource") // the underlying stream is managed by a subclass.
    254    @Override
    255    public void write(final byte[] b, final int off, final int len) throws IOException {
    256        IOUtils.checkFromIndexSize(b, off, len);
    257        // TODO we could write the sub-array up the threshold, fire the event,
    258        // and then write the rest so the event is always fired at the precise point.
    259        checkThreshold(len);
    260        // TODO for 4.0: Replace with getOutputStream()
    261        getStream().write(b, off, len);
    262        written += len;
    263    }
    264
    265    /**
    266     * Writes the specified byte to this output stream.
    267     *
    268     * @param b The byte to be written.
    269     * @throws IOException if an error occurs.
    270     */
    271    @SuppressWarnings("resource") // the underlying stream is managed by a subclass.
    272    @Override
    273    public void write(final int b) throws IOException {
    274        checkThreshold(1);
    275        // TODO for 4.0: Replace with getOutputStream()
    276        getStream().write(b);
    277        written++;
    278    }
    279}