samza-test/src/main/java/org/apache/samza/system/mock/MockSystemConsumer.java

/*
 * 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.samza.system.mock;

import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;

import org.apache.samza.metrics.MetricsRegistryMap;
import org.apache.samza.system.IncomingMessageEnvelope;
import org.apache.samza.system.SystemStreamPartition;
import org.apache.samza.util.BlockingEnvelopeMap;
import org.apache.samza.util.Clock;

/**
 * MockSystemConsumer is a class that simulates a multi-threaded consumer that
 * uses BlockingEnvelopeMap. The primary use for this class is to do performance
 * testing.
 *
 * This class works by starting up (threadCount) threads. Each thread adds
 * (messagesPerBatch) to the BlockingEnvelopeMap, then sleeps for
 * (brokerSleepMs). The sleep is important to simulate network latency when
 * executing a fetch against a remote streaming system (i.e. Kafka).
 */
public class MockSystemConsumer extends BlockingEnvelopeMap {
  private final int messagesPerBatch;
  private final int threadCount;
  private final int brokerSleepMs;

  /**
   * The SystemStreamPartitions that this consumer is in charge of.
   */
  private final Set<SystemStreamPartition> ssps;

  /**
   * The consumer threads that are putting IncomingMessageEnvelopes into
   * BlockingEnvelopeMap.
   */
  private List<Thread> threads;

  /**
   *
   * @param messagesPerBatch
   *          The number of messages to add to the BlockingEnvelopeMap before
   *          sleeping.
   * @param threadCount
   *          How many threads to run.
   * @param brokerSleepMs
   *          How long each thread should sleep between batch writes.
   */
  public MockSystemConsumer(int messagesPerBatch, int threadCount, int brokerSleepMs) {
    super(new MetricsRegistryMap("test-container-performance"), new Clock() {
      @Override
      public long currentTimeMillis() {
        return System.currentTimeMillis();
      }
    });

    this.messagesPerBatch = messagesPerBatch;
    this.threadCount = threadCount;
    this.brokerSleepMs = brokerSleepMs;
    this.ssps = new HashSet<SystemStreamPartition>();
    this.threads = new ArrayList<Thread>(threadCount);
  }

  /**
   * Assign SystemStreamPartitions to all of the threads, and start them up to
   * begin simulating consuming messages.
   */
  @Override
  public void start() {
    for (int i = 0; i < threadCount; ++i) {
      Set<SystemStreamPartition> threadSsps = new HashSet<SystemStreamPartition>();

      // Assign SystemStreamPartitions for this thread.
      for (SystemStreamPartition ssp : ssps) {
        if (Math.abs(ssp.hashCode()) % threadCount == i) {
          threadSsps.add(ssp);
        }
      }

      // Start thread.
      Thread thread = new Thread(new MockSystemConsumerRunnable(threadSsps), "MockSystemConsumer-" + i);
      thread.setDaemon(true);
      threads.add(thread);
      thread.start();
    }
  }

  /**
   * Kill all the threads, and shutdown.
   */
  @Override
  public void stop() {
    for (Thread thread : threads) {
      thread.interrupt();
    }

    try {
      for (Thread thread : threads) {
        thread.join();
      }
    } catch (InterruptedException e) {
    }
  }

  @Override
  public void register(SystemStreamPartition systemStreamPartition, String offset) {
    super.register(systemStreamPartition, offset);
    ssps.add(systemStreamPartition);
    setIsAtHead(systemStreamPartition, true);
  }

  /**
   * The worker thread for MockSystemConsumer that simulates reading messages
   * from a remote streaming system (i.e. Kafka), and writing them to the
   * BlockingEnvelopeMap.
   */
  public class MockSystemConsumerRunnable implements Runnable {
    private final Set<SystemStreamPartition> ssps;

    public MockSystemConsumerRunnable(Set<SystemStreamPartition> ssps) {
      this.ssps = ssps;
    }

    @Override
    public void run() {
      try {
        while (!Thread.interrupted() && ssps.size() > 0) {
          Set<SystemStreamPartition> sspsToFetch = new HashSet<SystemStreamPartition>();

          // Only fetch messages when there are no outstanding messages left.
          for (SystemStreamPartition ssp : ssps) {
            if (getNumMessagesInQueue(ssp) <= 0) {
              sspsToFetch.add(ssp);
            }
          }

          // Simulate a broker fetch request's network latency.
          Thread.sleep(brokerSleepMs);

          // Add messages to the BlockingEnvelopeMap.
          for (SystemStreamPartition ssp : sspsToFetch) {
            for (int i = 0; i < messagesPerBatch; ++i) {
              put(ssp, new IncomingMessageEnvelope(ssp, "0", "key", "value"));
            }
          }
        }
      } catch (InterruptedException e) {
        System.out.println("Got interrupt. Shutting down.");
      }
    }
  }
}