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 */
41 public class ThresholdingOutputStream extends OutputStream {
42
43 /**
44 * Noop output stream getter function.
45 */
46 private static final IOFunction<ThresholdingOutputStream, OutputStream> NOOP_OS_GETTER = os -> NullOutputStream.INSTANCE;
47
48 /**
49 * The threshold at which the event will be triggered.
50 */
51 private final int threshold;
52
53 /**
54 * Accepts reaching the threshold.
55 */
56 private final IOConsumer<ThresholdingOutputStream> thresholdConsumer;
57
58 /**
59 * Gets the output stream.
60 */
61 private final IOFunction<ThresholdingOutputStream, OutputStream> outputStreamGetter;
62
63 /**
64 * The number of bytes written to the output stream.
65 */
66 private long written;
67
68 /**
69 * Whether or not the configured threshold has been exceeded.
70 */
71 private boolean thresholdExceeded;
72
73 /**
74 * Constructs an instance of this class which will trigger an event at the specified threshold.
75 *
76 * @param threshold The number of bytes at which to trigger an event.
77 */
78 public ThresholdingOutputStream(final int threshold) {
79 this(threshold, IOConsumer.noop(), NOOP_OS_GETTER);
80 }
81
82 /**
83 * Constructs an instance of this class which will trigger an event at the specified threshold.
84 * A negative threshold has no meaning and will be treated as 0
85 *
86 * @param threshold The number of bytes at which to trigger an event.
87 * @param thresholdConsumer Accepts reaching the threshold.
88 * @param outputStreamGetter Gets the output stream.
89 * @since 2.9.0
90 */
91 public ThresholdingOutputStream(final int threshold, final IOConsumer<ThresholdingOutputStream> thresholdConsumer,
92 final IOFunction<ThresholdingOutputStream, OutputStream> outputStreamGetter) {
93 this.threshold = threshold < 0 ? 0 : threshold;
94 this.thresholdConsumer = thresholdConsumer == null ? IOConsumer.noop() : thresholdConsumer;
95 this.outputStreamGetter = outputStreamGetter == null ? NOOP_OS_GETTER : outputStreamGetter;
96 }
97
98 /**
99 * 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 }