View Javadoc
1   /*
2    * Licensed to the Apache Software Foundation (ASF) under one or more
3    * contributor license agreements.  See the NOTICE file distributed with
4    * this work for additional information regarding copyright ownership.
5    * The ASF licenses this file to You under the Apache License, Version 2.0
6    * (the "License"); you may not use this file except in compliance with
7    * the License.  You may obtain a copy of the License at
8    *
9    *      http://www.apache.org/licenses/LICENSE-2.0
10   *
11   * Unless required by applicable law or agreed to in writing, software
12   * distributed under the License is distributed on an "AS IS" BASIS,
13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14   * See the License for the specific language governing permissions and
15   * limitations under the License.
16   */
17  package org.apache.commons.io.output;
18  
19  import java.io.IOException;
20  import java.io.OutputStream;
21  
22  import org.apache.commons.io.function.IOConsumer;
23  import org.apache.commons.io.function.IOFunction;
24  
25  /**
26   * An output stream which triggers an event on the first write that causes
27   * the total number of bytes written to the stream to exceed a configured threshold,
28   * and every subsequent write. The event
29   * can be used, for example, to throw an exception if a maximum has been reached,
30   * or to switch the underlying stream when the threshold is exceeded.
31   *
32   * <p>
33   * This class overrides all {@link OutputStream} methods. However, these overrides ultimately call the corresponding
34   * methods in the underlying output stream implementation.
35   * </p>
36   * <p>
37   * NOTE: This implementation may trigger the event <em>before</em> the threshold is actually reached, since it triggers
38   * when a pending write operation would cause the threshold to be exceeded.
39   * </p>
40   * <p>
41   * See also the subclass {@link DeferredFileOutputStream}.
42   * </p>
43   *
44   * @see DeferredFileOutputStream
45   */
46  public class ThresholdingOutputStream extends OutputStream {
47  
48      /**
49       * Noop output stream getter function.
50       */
51      private static final IOFunction<ThresholdingOutputStream, OutputStream> NOOP_OS_GETTER = os -> NullOutputStream.INSTANCE;
52  
53      /**
54       * The threshold at which the event will be triggered.
55       */
56      private final int threshold;
57  
58      /**
59       * Accepts reaching the threshold.
60       */
61      private final IOConsumer<ThresholdingOutputStream> thresholdConsumer;
62  
63      /**
64       * Gets the output stream.
65       */
66      private final IOFunction<ThresholdingOutputStream, OutputStream> outputStreamGetter;
67  
68      /**
69       * The number of bytes written to the output stream.
70       */
71      private long written;
72  
73      /**
74       * Whether or not the configured threshold has been exceeded.
75       */
76      private boolean thresholdExceeded;
77  
78      /**
79       * Constructs an instance of this class which will trigger an event at the specified threshold.
80       *
81       * @param threshold The number of bytes at which to trigger an event.
82       */
83      public ThresholdingOutputStream(final int threshold) {
84          this(threshold, IOConsumer.noop(), NOOP_OS_GETTER);
85      }
86  
87      /**
88       * Constructs an instance of this class which will trigger an event at the specified threshold.
89       * A negative threshold has no meaning and will be treated as 0
90       *
91       * @param threshold The number of bytes at which to trigger an event.
92       * @param thresholdConsumer Accepts reaching the threshold.
93       * @param outputStreamGetter Gets the output stream.
94       * @since 2.9.0
95       */
96      public ThresholdingOutputStream(final int threshold, final IOConsumer<ThresholdingOutputStream> thresholdConsumer,
97          final IOFunction<ThresholdingOutputStream, OutputStream> outputStreamGetter) {
98          this.threshold = threshold < 0 ? 0 : threshold;
99          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 }