Source code

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