Skip to content

Commit 703b283

Browse files
garydgregorygarydgregory
committed
Add ThrottledInputStream
1 parent f7bd5f9 commit 703b283

4 files changed

Lines changed: 270 additions & 13 deletions

File tree

‎src/changes/changes.xml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -93,6 +93,7 @@ The <action> type attribute can be add,update,fix,remove.
9393
<action dev="ggregory" type="add" due-to="Gary Gregory">Add FileTimes.toUnixTime(FileTime).</action>
9494
<action dev="ggregory" type="add" due-to="Gary Gregory">Add BrokenInputStream.Builder.</action>
9595
<action dev="ggregory" type="add" due-to="Gary Gregory">Add PathUtils.getExtension(Path).</action>
96+
<action dev="ggregory" type="add" due-to="Gary Gregory">Add ThrottledInputStream.</action>
9697
<!-- UPDATE -->
9798
<action dev="ggregory" type="fix" due-to="Gary Gregory">Bump commons.bytebuddy.version from 1.14.10 to 1.14.11 #534.</action>
9899
</release>

‎src/main/java/org/apache/commons/io/input/ThrottledInputStream.java‎

Lines changed: 170 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,170 @@
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+
18+
package org.apache.commons.io.input;
19+
20+
import java.io.IOException;
21+
import java.io.InputStream;
22+
import java.io.InterruptedIOException;
23+
import java.time.Duration;
24+
import java.time.temporal.ChronoUnit;
25+
import java.util.concurrent.TimeUnit;
26+
27+
import org.apache.commons.io.build.AbstractStreamBuilder;
28+
29+
/**
30+
* Provides bandwidth throttling on a specified InputStream. It is implemented as a wrapper on top of another InputStream instance. The throttling works by
31+
* examining the number of bytes read from the underlying InputStream from the beginning, and sleep()ing for a time interval if the byte-transfer is found
32+
* exceed the specified tolerable maximum. (Thus, while the read-rate might exceed the maximum for a given short interval, the average tends towards the
33+
* specified maximum, overall.)
34+
* <p>
35+
* Inspired by Apache HBase's class of the same name.
36+
* </p>
37+
*
38+
* @since 2.16.0
39+
*/
40+
public final class ThrottledInputStream extends CountingInputStream {
41+
42+
/**
43+
* Builds a new {@link QueueInputStream} instance.
44+
* <h2>Using NIO</h2>
45+
*
46+
* <pre>{@code
47+
* ThrottledInputStream s = ThrottledInputStream.builder().setPath(Paths.get("MyFile.xml")).setMaxBytesPerSecond(100_000).get();
48+
* }
49+
* </pre>
50+
*
51+
* <h2>Using IO</h2>
52+
*
53+
* <pre>{@code
54+
* ThrottledInputStream s = ThrottledInputStream.builder().setFile(new File("MyFile.xml")).setMaxBytesPerSecond(100_000).get();
55+
* }
56+
* </pre>
57+
*
58+
* <pre>{@code
59+
* ThrottledInputStream s = ThrottledInputStream.builder().setInputStream(inputStream).setMaxBytesPerSecond(100_000).get();
60+
* }
61+
* </pre>
62+
*/
63+
public static class Builder extends AbstractStreamBuilder<ThrottledInputStream, Builder> {
64+
65+
/**
66+
* Effectively not throttled.
67+
*/
68+
private long maxBytesPerSecond = Long.MAX_VALUE;
69+
70+
@SuppressWarnings("resource")
71+
@Override
72+
public ThrottledInputStream get() throws IOException {
73+
return new ThrottledInputStream(getInputStream(), maxBytesPerSecond);
74+
}
75+
76+
/**
77+
* Sets the maximum bytes per second.
78+
*
79+
* @param maxBytesPerSecond the maximum bytes per second.
80+
*/
81+
public void setMaxBytesPerSecond(final long maxBytesPerSecond) {
82+
this.maxBytesPerSecond = maxBytesPerSecond;
83+
}
84+
85+
}
86+
87+
/**
88+
* Constructs a new {@link Builder}.
89+
*
90+
* @return a new {@link Builder}.
91+
*/
92+
public static Builder builder() {
93+
return new Builder();
94+
}
95+
96+
static long toSleepMillis(final long bytesRead, final long maxBytesPerSec, final long elapsedMillis) {
97+
assert elapsedMillis >= 0 : "The elapsed time should be greater or equal to zero";
98+
if (bytesRead <= 0 || maxBytesPerSec <= 0 || elapsedMillis == 0) {
99+
return 0;
100+
}
101+
// We use this class to load the single source file, so the bytesRead
102+
// and maxBytesPerSec aren't greater than Double.MAX_VALUE.
103+
// We can get the precise sleep time by using the double value.
104+
final long millis = (long) ((double) bytesRead / (double) maxBytesPerSec * 1000 - elapsedMillis);
105+
if (millis <= 0) {
106+
return 0;
107+
}
108+
return millis;
109+
}
110+
111+
private final long maxBytesPerSecond;
112+
private final long startTime = System.currentTimeMillis();
113+
private Duration totalSleepDuration = Duration.ZERO;
114+
115+
private ThrottledInputStream(final InputStream proxy, final long maxBytesPerSecond) {
116+
super(proxy);
117+
assert maxBytesPerSecond > 0 : "Bandwidth " + maxBytesPerSecond + " is invalid.";
118+
this.maxBytesPerSecond = maxBytesPerSecond;
119+
}
120+
121+
/**
122+
* Gets the read-rate from this stream, since creation. Calculated as bytesRead/elapsedTimeSinceStart.
123+
*
124+
* @return Read rate, in bytes/sec.
125+
*/
126+
public long getBytesPerSecond() {
127+
final long elapsedSeconds = (System.currentTimeMillis() - startTime) / 1000;
128+
if (elapsedSeconds == 0) {
129+
return getByteCount();
130+
}
131+
return getByteCount() / elapsedSeconds;
132+
}
133+
134+
private long getSleepMillis() {
135+
return toSleepMillis(getByteCount(), maxBytesPerSecond, System.currentTimeMillis() - startTime);
136+
}
137+
138+
/**
139+
* Gets the total duration spent in sleep.
140+
*
141+
* @return Duration spent in sleep.
142+
*/
143+
public Duration getTotalSleepDuration() {
144+
return totalSleepDuration;
145+
}
146+
147+
@Override
148+
protected void beforeRead(final int n) throws IOException {
149+
throttle();
150+
}
151+
152+
private void throttle() throws InterruptedIOException {
153+
final long sleepMillis = getSleepMillis();
154+
if (sleepMillis > 0) {
155+
totalSleepDuration = totalSleepDuration.plus(sleepMillis, ChronoUnit.MILLIS);
156+
try {
157+
TimeUnit.MILLISECONDS.sleep(sleepMillis);
158+
} catch (final InterruptedException e) {
159+
throw new InterruptedIOException("Thread aborted");
160+
}
161+
}
162+
}
163+
164+
/** {@inheritDoc} */
165+
@Override
166+
public String toString() {
167+
return "ThrottledInputStream[bytesRead=" + getCount() + ", maxBytesPerSec=" + maxBytesPerSecond + ", bytesPerSec=" + getBytesPerSecond()
168+
+ ", totalSleepDuration=" + totalSleepDuration + ']';
169+
}
170+
}

‎src/test/java/org/apache/commons/io/input/ProxyInputStreamTest.java‎

Lines changed: 28 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -29,8 +29,10 @@
2929

3030
/**
3131
* Tests {@link ProxyInputStream}.
32+
*
33+
* @param <T> The actual type tested.
3234
*/
33-
public class ProxyInputStreamTest {
35+
public class ProxyInputStreamTest<T extends ProxyInputStream> {
3436

3537
private static final class ProxyInputStreamFixture extends ProxyInputStream {
3638

@@ -40,10 +42,22 @@ public ProxyInputStreamFixture(final InputStream proxy) {
4042

4143
}
4244

45+
@SuppressWarnings({ "resource", "unused" }) // For subclasses
46+
protected T createFixture() throws IOException {
47+
return (T) new ProxyInputStreamFixture(createProxySource());
48+
}
49+
50+
protected InputStream createProxySource() {
51+
return CharSequenceInputStream.builder().setCharSequence("abc").get();
52+
}
53+
54+
protected void testEos(final T inputStream) {
55+
// empty
56+
}
57+
4358
@Test
4459
public void testRead() throws IOException {
45-
try (CharSequenceInputStream proxy = CharSequenceInputStream.builder().setCharSequence("abc").get();
46-
ProxyInputStream inputStream = new ProxyInputStreamFixture(proxy)) {
60+
try (T inputStream = createFixture()) {
4761
int found = inputStream.read();
4862
assertEquals('a', found);
4963
found = inputStream.read();
@@ -52,39 +66,39 @@ public void testRead() throws IOException {
5266
assertEquals('c', found);
5367
found = inputStream.read();
5468
assertEquals(-1, found);
69+
testEos(inputStream);
5570
}
5671
}
5772

5873
@Test
5974
public void testReadArrayAtMiddleFully() throws IOException {
60-
try (CharSequenceInputStream proxy = CharSequenceInputStream.builder().setCharSequence("abc").get();
61-
ProxyInputStream inputStream = new ProxyInputStreamFixture(proxy)) {
75+
try (T inputStream = createFixture()) {
6276
final byte[] dest = new byte[5];
6377
int found = inputStream.read(dest, 2, 3);
6478
assertEquals(3, found);
6579
assertArrayEquals(new byte[] { 0, 0, 'a', 'b', 'c' }, dest);
6680
found = inputStream.read(dest, 2, 3);
6781
assertEquals(-1, found);
82+
testEos(inputStream);
6883
}
6984
}
7085

7186
@Test
7287
public void testReadArrayAtStartFully() throws IOException {
73-
try (CharSequenceInputStream proxy = CharSequenceInputStream.builder().setCharSequence("abc").get();
74-
ProxyInputStream inputStream = new ProxyInputStreamFixture(proxy)) {
88+
try (T inputStream = createFixture()) {
7589
final byte[] dest = new byte[5];
7690
int found = inputStream.read(dest, 0, 5);
7791
assertEquals(3, found);
7892
assertArrayEquals(new byte[] { 'a', 'b', 'c', 0, 0 }, dest);
7993
found = inputStream.read(dest, 0, 5);
8094
assertEquals(-1, found);
95+
testEos(inputStream);
8196
}
8297
}
8398

8499
@Test
85100
public void testReadArrayAtStartPartial() throws IOException {
86-
try (CharSequenceInputStream proxy = CharSequenceInputStream.builder().setCharSequence("abc").get();
87-
ProxyInputStream inputStream = new ProxyInputStreamFixture(proxy)) {
101+
try (T inputStream = createFixture()) {
88102
final byte[] dest = new byte[5];
89103
int found = inputStream.read(dest, 0, 2);
90104
assertEquals(2, found);
@@ -95,26 +109,26 @@ public void testReadArrayAtStartPartial() throws IOException {
95109
assertArrayEquals(new byte[] { 'c', 0, 0, 0, 0 }, dest);
96110
found = inputStream.read(dest, 0, 2);
97111
assertEquals(-1, found);
112+
testEos(inputStream);
98113
}
99114
}
100115

101116
@Test
102117
public void testReadArrayFully() throws IOException {
103-
try (CharSequenceInputStream proxy = CharSequenceInputStream.builder().setCharSequence("abc").get();
104-
ProxyInputStream inputStream = new ProxyInputStreamFixture(proxy)) {
118+
try (T inputStream = createFixture()) {
105119
final byte[] dest = new byte[5];
106120
int found = inputStream.read(dest);
107121
assertEquals(3, found);
108122
assertArrayEquals(new byte[] { 'a', 'b', 'c', 0, 0 }, dest);
109123
found = inputStream.read(dest);
110124
assertEquals(-1, found);
125+
testEos(inputStream);
111126
}
112127
}
113128

114129
@Test
115130
public void testReadArrayPartial() throws IOException {
116-
try (CharSequenceInputStream proxy = CharSequenceInputStream.builder().setCharSequence("abc").get();
117-
ProxyInputStream inputStream = new ProxyInputStreamFixture(proxy)) {
131+
try (T inputStream = createFixture()) {
118132
final byte[] dest = new byte[2];
119133
int found = inputStream.read(dest);
120134
assertEquals(2, found);
@@ -125,6 +139,7 @@ public void testReadArrayPartial() throws IOException {
125139
assertArrayEquals(new byte[] { 'c', 0 }, dest);
126140
found = inputStream.read(dest);
127141
assertEquals(-1, found);
142+
testEos(inputStream);
128143
}
129144
}
130145

‎src/test/java/org/apache/commons/io/input/ThrottledInputStreamTest.java‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
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+
18+
package org.apache.commons.io.input;
19+
20+
import static org.junit.jupiter.api.Assertions.assertEquals;
21+
22+
import java.io.IOException;
23+
import java.time.Duration;
24+
25+
import org.junit.jupiter.api.Test;
26+
27+
/**
28+
* Tests {@link ThrottledInputStream}.
29+
*/
30+
public class ThrottledInputStreamTest extends ProxyInputStreamTest<ThrottledInputStream> {
31+
32+
@Override
33+
@SuppressWarnings("resource")
34+
protected ThrottledInputStream createFixture() throws IOException {
35+
return ThrottledInputStream.builder().setInputStream(createProxySource()).get();
36+
}
37+
38+
@Test
39+
public void testCalSleepTimeMs() {
40+
// case 0: initial - no read, no sleep
41+
assertEquals(0, ThrottledInputStream.toSleepMillis(0, 10_000, 1_000));
42+
43+
// case 1: no threshold
44+
assertEquals(0, ThrottledInputStream.toSleepMillis(Long.MAX_VALUE, 0, 1_000));
45+
assertEquals(0, ThrottledInputStream.toSleepMillis(Long.MAX_VALUE, -1, 1_000));
46+
47+
// case 2: too fast
48+
assertEquals(1500, ThrottledInputStream.toSleepMillis(5, 2, 1_000));
49+
assertEquals(500, ThrottledInputStream.toSleepMillis(5, 2, 2_000));
50+
assertEquals(6500, ThrottledInputStream.toSleepMillis(15, 2, 1_000));
51+
52+
// case 3: too slow
53+
assertEquals(0, ThrottledInputStream.toSleepMillis(1, 2, 1_000));
54+
assertEquals(0, ThrottledInputStream.toSleepMillis(2, 2, 2_000));
55+
assertEquals(0, ThrottledInputStream.toSleepMillis(1, 2, 1_000));
56+
}
57+
58+
@Override
59+
protected void testEos(final ThrottledInputStream inputStream) {
60+
assertEquals(3, inputStream.getByteCount());
61+
}
62+
63+
@Test
64+
public void testGet() throws IOException {
65+
try (ThrottledInputStream inputStream = createFixture()) {
66+
inputStream.read();
67+
assertEquals(Duration.ZERO, inputStream.getTotalSleepDuration());
68+
}
69+
}
70+
71+
}

0 commit comments

Comments
 (0)