com.datastax.oss.driver.api.testinfra.simulacron.QueryCounter Maven / Gradle / Ivy

Go to download
/*
 * Copyright DataStax, 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 com.datastax.oss.driver.api.testinfra.simulacron;

import static org.assertj.core.api.Assertions.assertThat;
import static org.awaitility.Awaitility.await;

import com.datastax.oss.simulacron.common.cluster.QueryLog;
import com.datastax.oss.simulacron.server.BoundNode;
import com.datastax.oss.simulacron.server.BoundTopic;
import com.datastax.oss.simulacron.server.listener.QueryListener;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.Predicate;

/**
 * A convenience utility that keeps track of the number of queries matching a filter received by
 * {@link BoundNode}s.
 *
 * One tricky thing about validating query counts in context of testing is in cases where you
 * don't require waiting for a node to respond. In this case it's possible that a user would check
 * count criteria before the node has even processed the message. This class offers the capability
 * to wait a specified amount of time when asserting query arrival counts on nodes.
 */
public class QueryCounter {

  private final long beforeTimeout;
  private final TimeUnit beforeUnit;
  private final AtomicInteger totalCount = new AtomicInteger(0);
  private final ConcurrentHashMap countMap = new ConcurrentHashMap<>();

  public enum NotificationMode {
    BEFORE_PROCESSING,
    AFTER_PROCESSING
  }

  private QueryCounter(
      BoundTopic topic,
      NotificationMode notificationMode,
      Predicate queryLogFilter,
      long beforeTimeout,
      TimeUnit beforeUnit) {
    this.beforeTimeout = beforeTimeout;
    this.beforeUnit = beforeUnit;
    QueryListener listener =
        (boundNode, queryLog) -> {
          totalCount.incrementAndGet();
          countMap.merge(boundNode.getId().intValue(), 1, Integer::sum);
        };
    topic.registerQueryListener(
        listener, notificationMode == NotificationMode.AFTER_PROCESSING, queryLogFilter);
  }

  /** Creates a builder that tracks counts for the given {@link BoundTopic} (cluster, dc, node). */
  public static QueryCounterBuilder builder(BoundTopic topic) {
    return new QueryCounterBuilder(topic);
  }

  /** Clears all counters. */
  public void clearCounts() {
    totalCount.set(0);
    countMap.clear();
  }

  /**
   * Asserts that the total number of requests received matching filter criteria matches the
   * expected count within the configured time period.
   */
  public void assertTotalCount(int expected) {
    await()
        .pollInterval(10, TimeUnit.MILLISECONDS)
        .atMost(beforeTimeout, beforeUnit)
        .untilAsserted(() -> assertThat(totalCount.get()).isEqualTo(expected));
  }

  /**
   * Asserts that the total number of requests received matcher filter criteria matches the expected
   * count for each node within the configured time period.
   *
   * @param counts The expected node counts, with the value at each index matching the expected
   *     count for that node id (i.e. index 0 = node id 0 expected count).
   */
  public void assertNodeCounts(int... counts) {
    Map expectedCounts = new HashMap<>();
    for (int id = 0; id < counts.length; id++) {
      int count = counts[id];
      if (count > 0) {
        expectedCounts.put(id, counts[id]);
      }
    }
    await()
        .pollInterval(10, TimeUnit.MILLISECONDS)
        .atMost(beforeTimeout, beforeUnit)
        .untilAsserted(() -> assertThat(countMap).containsAllEntriesOf(expectedCounts));
  }

  public static class QueryCounterBuilder {

    @SuppressWarnings("deprecation")
    private static Predicate DEFAULT_FILTER = (q) -> !q.getQuery().isEmpty();

    private Predicate queryLogFilter = DEFAULT_FILTER;
    private BoundTopic topic;
    private NotificationMode notificationMode = NotificationMode.BEFORE_PROCESSING;
    private long beforeTimeout = 1;
    private TimeUnit beforeUnit = TimeUnit.SECONDS;

    private QueryCounterBuilder(BoundTopic topic) {
      this.topic = topic;
    }

    /**
     * The filter to apply to consider a message received by the node, if not provided we consider
     * all messages that are queries.
     */
    public QueryCounterBuilder withFilter(Predicate queryLogFilter) {
      this.queryLogFilter = queryLogFilter;
      return this;
    }

    /** Whether or not simulacron should notify before or after the message is processed. */
    public QueryCounterBuilder withNotification(NotificationMode notificationMode) {
      this.notificationMode = notificationMode;
      return this;
    }

    /**
     * Up to how long we check counts to match. If counts don't match after this time, an exception
     * is thrown.
     */
    public QueryCounterBuilder before(long timeout, TimeUnit unit) {
      this.beforeTimeout = timeout;
      this.beforeUnit = unit;
      return this;
    }

    public QueryCounter build() {
      return new QueryCounter(topic, notificationMode, queryLogFilter, beforeTimeout, beforeUnit);
    }
  }
}