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