content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you 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 org.apache.cassandra.transport;
  
  import java.util.Collection;
  import java.util.Map;
  import java.util.concurrent.CountDownLatch;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicLong;
  import java.util.concurrent.atomic.AtomicReference;
  import java.util.stream.Collectors;
  
  import com.codahale.metrics.Meter;
  import com.google.common.base.Ticker;
  
  import org.awaitility.Awaitility;
  
  import org.junit.Before;
  import org.junit.BeforeClass;
  import org.junit.Test;
  import org.junit.runner.RunWith;
  import org.junit.runners.Parameterized;
  
  import org.apache.cassandra.config.DatabaseDescriptor;
  import org.apache.cassandra.exceptions.OverloadedException;
  import org.apache.cassandra.metrics.CassandraMetricsRegistry;
  import org.apache.cassandra.service.StorageService;
  import org.apache.cassandra.utils.Throwables;
  
  import static org.junit.Assert.assertEquals;
  import static org.junit.Assert.assertNull;
  import static org.junit.Assert.assertTrue;
  import static org.junit.Assert.fail;
  
  import static org.apache.cassandra.Util.spinAssertEquals;
  import static org.apache.cassandra.transport.ProtocolVersion.V4;
  
  @SuppressWarnings("UnstableApiUsage")
  @RunWith(Parameterized.class)
  public class RateLimitingTest extends NativeProtocolLimitsTestBase
  {
      public static final String BACKPRESSURE_WARNING_SNIPPET = "Request breached global limit";
      
      private static final int LARGE_PAYLOAD_THRESHOLD_BYTES = 1000;
      private static final int OVERLOAD_PERMITS_PER_SECOND = 1;
  
      private static final long MAX_LONG_CONFIG_VALUE = Long.MAX_VALUE - 1;
  
      @Parameterized.Parameters(name="{0}")
      public static Collection<Object[]> versions()
      {
          return ProtocolVersion.SUPPORTED.stream()
                                          .map(v -> new Object[]{v})
                                          .collect(Collectors.toList());
      }
  
      public RateLimitingTest(ProtocolVersion version)
      {
          super(version);
      }
  
      private AtomicLong tick;
      private Ticker ticker;
  
      @BeforeClass
      public static void setup()
      {
          // If we don't exceed the queue capacity, we won't actually use the global/endpoint 
          // bytes-in-flight limits, and the assertions we make below around releasing them would be useless.
          DatabaseDescriptor.setNativeTransportReceiveQueueCapacityInBytes(1);
  
          requireNetwork();
      }
  
      @Before
      public void resetLimits()
      {
          // Reset to the original start time in case a test advances the clock.
          tick = new AtomicLong(ClientResourceLimits.GLOBAL_REQUEST_LIMITER.getStartedNanos());
  
          ticker = new Ticker()
          {
              @Override
              public long read()
              {
                  return tick.get();
              }
          };
  
          ClientResourceLimits.setGlobalLimit(MAX_LONG_CONFIG_VALUE);
      }
  
  
      @Test
      public void shouldThrowOnOverloadSmallMessages() throws Exception
      {
          int payloadSize = LARGE_PAYLOAD_THRESHOLD_BYTES / 4;
          testOverload(payloadSize, true);
      }
  
      @Test
      public void shouldThrowOnOverloadLargeMessages() throws Exception
      {
          int payloadSize = LARGE_PAYLOAD_THRESHOLD_BYTES * 2;
          testOverload(payloadSize, true);
      }
  
      @Test
      public void shouldBackpressureSmallMessages() throws Exception
      {
          int payloadSize = LARGE_PAYLOAD_THRESHOLD_BYTES / 4;
          testOverload(payloadSize, false);
      }
  
      @Test
      public void shouldBackpressureLargeMessages() throws Exception
      {
          int payloadSize = LARGE_PAYLOAD_THRESHOLD_BYTES * 2;
          testOverload(payloadSize, false);
      }
  
      @Test
      public void shouldReleaseSmallMessageOnBytesInFlightOverload() throws Exception
      {
          testBytesInFlightOverload(LARGE_PAYLOAD_THRESHOLD_BYTES / 4);
      }
  
      @Test
      public void shouldReleaseLargeMessageOnBytesInFlightOverload() throws Exception
      {
          testBytesInFlightOverload(LARGE_PAYLOAD_THRESHOLD_BYTES * 2);
      }
  
      private void testBytesInFlightOverload(int payloadSize) throws Exception
      {
          int emulatedConcurrentMessageSize = payloadSize * 3 / 2;
          try (SimpleClient client = client(true, LARGE_PAYLOAD_THRESHOLD_BYTES))
          {
              StorageService.instance.setNativeTransportRateLimitingEnabled(false);
              createTable(client);
  
              StorageService.instance.setNativeTransportRateLimitingEnabled(true);
              ClientResourceLimits.GLOBAL_REQUEST_LIMITER.setRate(OVERLOAD_PERMITS_PER_SECOND, ticker);
              // test message = 1x
              // emulated concurrent message = 1.5x
              // test message + emulated concurrent message = 2.5x > 2x set as a global limit
              ClientResourceLimits.setGlobalLimit(payloadSize * 2);
              emulateInFlightConcurrentMessage(emulatedConcurrentMessageSize);
  
              try
              {
                  // The first query takes the one available permit, but should fail on the bytes in flight limit.
                  client.execute(queryMessage(payloadSize));
              }
              catch (RuntimeException e)
              {
                  assertTrue(Throwables.anyCauseMatches(e, cause -> cause instanceof OverloadedException));
              }
          }
          finally
          {
              // Sanity check bytes in flight limiter.
              Awaitility.await().untilAsserted(() -> assertEquals(emulatedConcurrentMessageSize, ClientResourceLimits.getCurrentGlobalUsage()));
              StorageService.instance.setNativeTransportRateLimitingEnabled(false);
          }
      }
  
      private void testOverload(int payloadSize, boolean throwOnOverload) throws Exception
      {
          try (SimpleClient client = client(throwOnOverload, LARGE_PAYLOAD_THRESHOLD_BYTES))
          {
              StorageService.instance.setNativeTransportRateLimitingEnabled(false);
              createTable(client);
  
              StorageService.instance.setNativeTransportRateLimitingEnabled(true);
              ClientResourceLimits.GLOBAL_REQUEST_LIMITER.setRate(OVERLOAD_PERMITS_PER_SECOND, ticker);
  
              if (throwOnOverload)
                  testThrowOnOverload(payloadSize, client);
              else
              {
                  testBackpressureOnOverload(payloadSize, client);
              }   
          }
          finally
          {
              // Sanity the check bytes in flight limiter.
              Awaitility.await().untilAsserted(() -> assertEquals(0, ClientResourceLimits.getCurrentGlobalUsage()));
              StorageService.instance.setNativeTransportRateLimitingEnabled(false);
          }
      }
  
      private void testBackpressureOnOverload(int payloadSize, SimpleClient client) throws Exception
      {
          // The first query takes the one available permit.
          Message.Response firstResponse = client.execute(queryMessage(payloadSize));
          assertEquals(0, getPausedConnectionsGauge().getValue().intValue());
          assertNoWarningContains(firstResponse, BACKPRESSURE_WARNING_SNIPPET);
          
          // The second query activates backpressure.
          long overloadQueryStartTime = System.currentTimeMillis();
          Message.Response response = client.execute(queryMessage(payloadSize));
  
          // V3 does not support client warnings, but otherwise we should get one for this query.
          if (version.isGreaterOrEqualTo(V4))
              assertWarningsContain(response, BACKPRESSURE_WARNING_SNIPPET);
  
          AtomicReference<Throwable> error = new AtomicReference<>();
          CountDownLatch started = new CountDownLatch(1);
          CountDownLatch complete = new CountDownLatch(1);
          AtomicReference<Message.Response> pausedQueryResponse = new AtomicReference<>();
          
          Thread queryRunner = new Thread(() ->
          {
              try
              {
                  started.countDown();
                  pausedQueryResponse.set(client.execute(queryMessage(payloadSize)));
                  complete.countDown();
              }
              catch (Throwable t)
              {
                  error.set(t);
              }
          });
  
          // Advance the rater limiter so that this query will see an available permit. This also
          // means it should not produce a client warning, which we verify below.
          // (Note that we advance 2 intervals for the 2 prior queries.)
          tick.addAndGet(2 * ClientResourceLimits.GLOBAL_REQUEST_LIMITER.getIntervalNanos());
          
          queryRunner.start();
  
          // ...and the request should complete without error.
          assertTrue(complete.await(SimpleClient.TIMEOUT_SECONDS + 1, TimeUnit.SECONDS));
          assertNull(error.get());
          assertNoWarningContains(pausedQueryResponse.get(), BACKPRESSURE_WARNING_SNIPPET);
  
          // At least the number of milliseconds in the permit interval should already have elapsed 
          // since the start of the query that pauses the connection.
          double permitIntervalMillis = (double) TimeUnit.SECONDS.toMillis(1L) / OVERLOAD_PERMITS_PER_SECOND;
          long sinceQueryStarted = System.currentTimeMillis() - overloadQueryStartTime;
          long remainingMillis = ((long) permitIntervalMillis) - sinceQueryStarted;
          assertTrue("Query completed before connection unpause!", remainingMillis <= 0);
          
          spinAssertEquals("Timed out after waiting 5 seconds for paused connections metric to normalize.",
                           0, () -> getPausedConnectionsGauge().getValue(), 5, TimeUnit.SECONDS);
      }
  
      private void testThrowOnOverload(int payloadSize, SimpleClient client)
      {
          // The first query takes the one available permit...
          long dispatchedPrior = getRequestDispatchedMeter().getCount();
          client.execute(queryMessage(payloadSize));
          assertEquals(dispatchedPrior + 1, getRequestDispatchedMeter().getCount());
          
          try
          {   
              // ...and the second breaches the limit....
              client.execute(queryMessage(payloadSize));
          }
          catch (RuntimeException e)
          {
              assertTrue(Throwables.anyCauseMatches(e, cause -> cause instanceof OverloadedException));
          }
  
          // The last request attempt was rejected and therefore not dispatched.
          assertEquals(dispatchedPrior + 1, getRequestDispatchedMeter().getCount());
  
          // Advance the timeline and verify that we can take a permit again.
          // (Note that we don't take one when we throw on overload.)
          tick.addAndGet(ClientResourceLimits.GLOBAL_REQUEST_LIMITER.getIntervalNanos());
          client.execute(queryMessage(payloadSize));
  
          assertEquals(dispatchedPrior + 2, getRequestDispatchedMeter().getCount());
      }
  
      protected static Meter getRequestDispatchedMeter()
      {
          String metricName = "org.apache.cassandra.metrics.Client.RequestDispatched";
          Map<String, Meter> metrics = CassandraMetricsRegistry.Metrics.getMeters((name, metric) -> name.equals(metricName));
          if (metrics.size() != 1)
              fail(String.format("Expected a single registered metric for request dispatched, found %s",metrics.size()));
          return metrics.get(metricName);
      }
  }
  