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