src/test/java/com/google/devtools/build/lib/remote/grpc/TokenBucketTest.java

// Copyright 2021 The Bazel Authors. All rights reserved.
//
// 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 com.google.devtools.build.lib.remote.grpc;

import static com.google.common.truth.Truth.assertThat;

import com.google.common.collect.ImmutableList;
import io.reactivex.rxjava3.core.Single;
import io.reactivex.rxjava3.observers.TestObserver;
import java.io.IOException;
import org.junit.Test;
import org.junit.runner.RunWith;
import org.junit.runners.JUnit4;

/** Tests for {@link TokenBucket} */
@RunWith(JUnit4.class)
public class TokenBucketTest {

  @Test
  public void acquireToken_smoke() {
    TokenBucket<Integer> bucket = new TokenBucket<>();
    assertThat(bucket.size()).isEqualTo(0);
    bucket.addToken(0);
    assertThat(bucket.size()).isEqualTo(1);

    TestObserver<Integer> observer = bucket.acquireToken().test();

    observer.assertValue(0).assertComplete();
    assertThat(bucket.size()).isEqualTo(0);
  }

  @Test
  public void acquireToken_releaseInitialTokens() {
    TokenBucket<Integer> bucket = new TokenBucket<>(ImmutableList.of(0));
    assertThat(bucket.size()).isEqualTo(1);

    TestObserver<Integer> observer = bucket.acquireToken().test();

    observer.assertValue(0).assertComplete();
    assertThat(bucket.size()).isEqualTo(0);
  }

  @Test
  public void acquireToken_multipleInitialTokens_releaseFirstToken() {
    TokenBucket<Integer> bucket = new TokenBucket<>(ImmutableList.of(0, 1));
    assertThat(bucket.size()).isEqualTo(2);

    TestObserver<Integer> observer = bucket.acquireToken().test();

    observer.assertValue(0).assertComplete();
    assertThat(bucket.size()).isEqualTo(1);
  }

  @Test
  public void acquireToken_multipleInitialTokens_releaseSecondToken() {
    TokenBucket<Integer> bucket = new TokenBucket<>(ImmutableList.of(0, 1));
    assertThat(bucket.size()).isEqualTo(2);
    bucket.acquireToken().test().assertValue(0).assertComplete();

    TestObserver<Integer> observer = bucket.acquireToken().test();

    observer.assertValue(1).assertComplete();
    assertThat(bucket.size()).isEqualTo(0);
  }

  @Test
  public void acquireToken_releaseTokenToPreviousObserver() {
    TokenBucket<Integer> bucket = new TokenBucket<>();
    TestObserver<Integer> observer = bucket.acquireToken().test();
    observer.assertEmpty();

    bucket.addToken(0);

    observer.assertValue(0).assertComplete();
    assertThat(bucket.size()).isEqualTo(0);
  }

  @Test
  public void acquireToken_notReleaseTokenToDisposedObserver() {
    TokenBucket<Integer> bucket = new TokenBucket<>();
    TestObserver<Integer> observer = bucket.acquireToken().test();

    observer.dispose();
    bucket.addToken(0);

    observer.assertEmpty();
    assertThat(bucket.size()).isEqualTo(1);
  }

  @Test
  public void acquireToken_disposeAfterTokenAcquired() {
    TokenBucket<Integer> bucket = new TokenBucket<>();
    TestObserver<Integer> observer = bucket.acquireToken().test();

    bucket.addToken(0);
    bucket.addToken(1);

    observer.assertValue(0).assertComplete();
    assertThat(bucket.size()).isEqualTo(1);
  }

  @Test
  public void acquireToken_multipleObservers_onlyOneCanAcquire() {
    TokenBucket<Integer> bucket = new TokenBucket<>();
    TestObserver<Integer> observer1 = bucket.acquireToken().test();
    TestObserver<Integer> observer2 = bucket.acquireToken().test();

    bucket.addToken(0);

    if (!observer1.values().isEmpty()) {
      observer1.assertValue(0).assertComplete();
      observer2.assertEmpty();

      bucket.addToken(1);
      observer2.assertValue(1).assertComplete();
    } else {
      observer1.assertEmpty();
      observer2.assertValue(0).assertComplete();

      bucket.addToken(1);
      observer1.assertValue(1).assertComplete();
    }
  }

  @Test
  public void acquireToken_reSubscription_waitAvailableToken() {
    TokenBucket<Integer> bucket = new TokenBucket<>();
    bucket.addToken(0);
    Single<Integer> tokenSingle = bucket.acquireToken();

    TestObserver<Integer> observer1 = tokenSingle.test();
    TestObserver<Integer> observer2 = tokenSingle.test();

    observer1.assertValue(0).assertComplete();
    observer2.assertEmpty();
  }

  @Test
  public void acquireToken_reSubscription_acquireNewToken() {
    TokenBucket<Integer> bucket = new TokenBucket<>();
    bucket.addToken(0);
    Single<Integer> tokenSingle = bucket.acquireToken();
    TestObserver<Integer> observer1 = tokenSingle.test();
    TestObserver<Integer> observer2 = tokenSingle.test();

    bucket.addToken(1);

    observer1.assertValue(0).assertComplete();
    observer2.assertValue(1).assertComplete();
  }

  @Test
  public void acquireToken_reSubscription_acquireNextToken() {
    TokenBucket<Integer> bucket = new TokenBucket<>();
    bucket.addToken(0);
    bucket.addToken(1);
    Single<Integer> tokenSingle = bucket.acquireToken();

    TestObserver<Integer> observer1 = tokenSingle.test();
    TestObserver<Integer> observer2 = tokenSingle.test();

    observer1.assertValue(0).assertComplete();
    observer2.assertValue(1).assertComplete();
  }

  @Test
  public void acquireToken_disposed_tokenRemains() {
    TokenBucket<Integer> bucket = new TokenBucket<>();
    TestObserver<Integer> observer = bucket.acquireToken().test();
    observer.assertEmpty();

    observer.dispose();
    bucket.addToken(0);

    assertThat(bucket.size()).isEqualTo(1);
  }

  @Test
  public void close_errorAfterClose() throws IOException {
    TokenBucket<Integer> bucket = new TokenBucket<>();
    bucket.addToken(0);
    bucket.close();

    TestObserver<Integer> observer = bucket.acquireToken().test();

    observer.assertError(
        e -> e instanceof IllegalStateException && e.getMessage().contains("closed"));
  }

  @Test
  public void close_errorPreviousObservers() throws IOException {
    TokenBucket<Integer> bucket = new TokenBucket<>();
    TestObserver<Integer> observer = bucket.acquireToken().test();

    bucket.close();

    observer.assertError(
        e -> e instanceof IllegalStateException && e.getMessage().contains("closed"));
  }
}