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 * https://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 }
275