aboutsummaryrefslogtreecommitdiff
path: root/okio/src/jvmTest/kotlin/okio/ThrottlerTakeTest.kt
blob: aab94b303095519a5d412a1bcc4a7ddf34506224 (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
/*
 * Copyright (C) 2018 Square, Inc.
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package okio

import org.assertj.core.api.Assertions.assertThat
import org.junit.Test
import java.util.concurrent.TimeUnit

class ThrottlerTakeTest {
  private var nowNanos = 0L
  private var elapsedNanos = 0L
  private val throttler = Throttler(allocatedUntil = nowNanos)

  @Test fun takeByByteCount() {
    throttler.bytesPerSecond(bytesPerSecond = 20, waitByteCount = 5, maxByteCount = 10)

    // We get the first 10 bytes immediately (that's maxByteCount).
    assertThat(take(100L)).isEqualTo(10L)
    assertElapsed(0L)

    // Wait a quarter second for each subsequent 5 bytes (that's waitByteCount).
    assertThat(take(100L)).isEqualTo(5L)
    assertElapsed(250L)

    assertThat(take(100L)).isEqualTo(5L)
    assertElapsed(250L)

    // Wait three quarters of a second to build up 15 bytes of potential.
    // Since maxByteCount = 10, there will only be 10 bytes of potential.
    sleep(750L)
    assertElapsed(750L)

    // We get 10 bytes immediately (that's maxByteCount again).
    assertThat(take(100L)).isEqualTo(10L)
    assertElapsed(0L)

    // Wait a quarter second for each subsequent 5 bytes (that's waitByteCount again).
    assertThat(take(100L)).isEqualTo(5L)
    assertElapsed(250L)
  }

  @Test fun takeFullyTimeElapsed() {
    throttler.bytesPerSecond(bytesPerSecond = 20, waitByteCount = 5, maxByteCount = 10)

    // We write the first 10 bytes immediately (that's maxByteCount again).
    takeFully(10L)
    assertElapsed(0L)

    // Wait a quarter second for each subsequent 5 bytes (that's waitByteCount).
    takeFully(5L)
    assertElapsed(250L)

    // Wait a half second for 10 bytes.
    takeFully(10L)
    assertElapsed(500L)

    // Wait a three quarters of a second to build up 15 bytes of potential.
    // Since maxByteCount = 10, there will only be 10 bytes of potential.
    sleep(750L)
    assertElapsed(750L)

    // We write the first 10 bytes immediately (that's maxByteCount again).
    // Wait a quarter second for each subsequent 5 bytes (that's waitByteCount again).
    takeFully(15L)
    assertElapsed(250L)
  }

  @Test fun takeFullyWhenSaturated() {
    throttler.bytesPerSecond(400L, 5L, 10L)

    // Saturate the throttler.
    assertThat(take(10L)).isEqualTo(10L)
    assertElapsed(0L)

    // At 400 bytes per second it takes 250 ms to read 100 bytes.
    takeFully(100L)
    assertElapsed(250L)
  }

  @Test fun takeFullyNoLimit() {
    throttler.bytesPerSecond(0L, 5L, 10L)
    takeFully(100L)
    assertElapsed(0L)
  }

  /**
   * We had a bug where integer division truncation would cause us to call wait() for 0 nanos. We
   * fixed it by minimizing integer division generally, and by handling that case specifically.
   */
  @Test fun infiniteWait() {
    throttler.bytesPerSecond(3, maxByteCount = 4, waitByteCount = 4)
    takeFully(7)
    assertElapsed(1000L)
  }

  /** Take at least the minimum and up to `byteCount` bytes, sleeping once if necessary. */
  private fun take(byteCount: Long): Long {
    val byteCountOrWaitNanos = throttler.byteCountOrWaitNanos(nowNanos, byteCount)
    if (byteCountOrWaitNanos >= 0L) return byteCountOrWaitNanos

    nowNanos += -byteCountOrWaitNanos

    val resultAfterWait = throttler.byteCountOrWaitNanos(nowNanos, byteCount)
    assertThat(resultAfterWait).isGreaterThan(0L)
    return resultAfterWait
  }

  /** Take all of `byteCount` bytes, advancing the clock until they're all taken. */
  private fun takeFully(byteCount: Long) {
    var remaining = byteCount
    while (remaining > 0L) {
      val byteCountOrWaitNanos = throttler.byteCountOrWaitNanos(nowNanos, remaining)
      if (byteCountOrWaitNanos >= 0L) {
        remaining -= byteCountOrWaitNanos
      } else {
        nowNanos += -byteCountOrWaitNanos
      }
    }
  }

  private fun assertElapsed(millis: Long) {
    elapsedNanos += TimeUnit.MILLISECONDS.toNanos(millis)
    assertThat(nowNanos).isEqualTo(elapsedNanos)
  }

  private fun sleep(millis: Long) {
    nowNanos += TimeUnit.MILLISECONDS.toNanos(millis)
  }
}