MetricsManagerServerTest.java

// Copyright 2016 Twitter. 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.twitter.heron.metricsmgr;

import java.io.IOException;

import com.google.protobuf.Message;

import org.junit.After;
import org.junit.Assert;
import org.junit.Before;
import org.junit.Test;

import com.twitter.heron.api.metric.MultiCountMetric;
import com.twitter.heron.common.basics.Communicator;
import com.twitter.heron.common.basics.NIOLooper;
import com.twitter.heron.common.basics.SysUtils;
import com.twitter.heron.common.basics.TestCommunicator;
import com.twitter.heron.common.network.HeronClient;
import com.twitter.heron.common.network.HeronServerTester;
import com.twitter.heron.common.network.StatusCode;
import com.twitter.heron.proto.system.Metrics;
import com.twitter.heron.spi.metricsmgr.metrics.ExceptionInfo;
import com.twitter.heron.spi.metricsmgr.metrics.MetricsInfo;
import com.twitter.heron.spi.metricsmgr.metrics.MetricsRecord;

import static com.twitter.heron.common.network.HeronServerTester.RESPONSE_RECEIVED_TIMEOUT;
import static org.junit.Assert.fail;

/**
 * MetricsManagerServer Tester.
 */
public class MetricsManagerServerTest {
  private static final String METRIC_NAME = "metric-name";
  private static final String METRIC_VALUE = "metric-value";

  private static final int METRICS_COUNT = 20;
  private static final int MESSAGE_SIZE = 10;
  private static final String STACK_TRACE = "stackTrace";
  private static final String LAST_TIME = "lastTime";
  private static final String FIRST_TIME = "firstTime";
  private static final String LOGGING = "logging";
  private static final int EXCEPTION_COUNT = 20;

  private MetricsManagerServer metricsManagerServer;
  private HeronServerTester serverTester;

  @Before
  public void before() throws IOException {
    metricsManagerServer = new MetricsManagerServer(new NIOLooper(), HeronServerTester.SERVER_HOST,
        SysUtils.getFreePort(), HeronServerTester.TEST_SOCKET_OPTIONS, new MultiCountMetric());

    serverTester = new HeronServerTester(metricsManagerServer,
        new MetricsManagerClientRequestHandler(),
        new HeronServerTester.SuccessResponseHandler(Metrics.MetricPublisherRegisterResponse.class,
            new MetricsManagerClientResponseHandler(MESSAGE_SIZE)), RESPONSE_RECEIVED_TIMEOUT);
  }

  @After
  public void after() {
    serverTester.stop();
  }

  /**
   * Method: addSinkCommunicator(Communicator<MetricsRecord> communicator)
   */
  @Test
  public void testAddSinkCommunicator() {
    Communicator<MetricsRecord> sinkCommunicator = new Communicator<>();
    metricsManagerServer.addSinkCommunicator(sinkCommunicator);
    Assert.assertTrue(metricsManagerServer.removeSinkCommunicator(sinkCommunicator));
  }

  /**
   * Method: removeSinkCommunicator(Communicator<MetricsRecord> communicator)
   */
  @Test
  public void testRemoveSinkCommunicator() {
    Communicator<MetricsRecord> sinkCommunicator = new Communicator<>();
    metricsManagerServer.addSinkCommunicator(sinkCommunicator);
    Assert.assertTrue(metricsManagerServer.removeSinkCommunicator(sinkCommunicator));
  }

  /**
   * Method: addSinkCommunicator(Communicator<MetricsRecord> communicator)
   */
  @Test
  public void testMetricsManagerServer() throws InterruptedException {
    TestCommunicator<MetricsRecord> sinkCommunicator =
        new TestCommunicator<>(MESSAGE_SIZE);
    metricsManagerServer.addSinkCommunicator(sinkCommunicator);

    serverTester.start();

    try {
      sinkCommunicator.awaitOffers(RESPONSE_RECEIVED_TIMEOUT);
    } catch (InterruptedException e) {
      fail(String.format("awaitOffers failed to release until timeout of %s was reached.",
          RESPONSE_RECEIVED_TIMEOUT));
    }

    int messages = 0;
    while (!sinkCommunicator.isEmpty()) {
      int exceptions = 0;
      int metrics = 0;

      MetricsRecord record = sinkCommunicator.poll();

      Assert.assertEquals("hostname:0/component/instance-id", record.getSource());
      Assert.assertEquals("default", record.getContext());

      for (MetricsInfo info : record.getMetrics()) {
        Assert.assertEquals(METRIC_NAME, info.getName());
        Assert.assertEquals(METRIC_VALUE, info.getValue());
        metrics++;
      }
      Assert.assertEquals(METRICS_COUNT, metrics);

      for (ExceptionInfo info : record.getExceptions()) {
        Assert.assertEquals(STACK_TRACE, info.getStackTrace());
        Assert.assertEquals(LAST_TIME, info.getLastTime());
        Assert.assertEquals(FIRST_TIME, info.getFirstTime());
        Assert.assertEquals(LOGGING, info.getLogging());
        Assert.assertEquals(EXCEPTION_COUNT, info.getCount());
        exceptions++;
      }
      Assert.assertEquals(METRICS_COUNT, exceptions);

      messages++;
    }
    Assert.assertEquals(MESSAGE_SIZE, messages);
  }

  private class MetricsManagerClientRequestHandler implements HeronServerTester.TestRequestHandler {

    @Override
    public Message getRequestMessage() {
      Metrics.MetricPublisher publisher = Metrics.MetricPublisher.newBuilder().
          setHostname("hostname").
          setPort(0).
          setComponentName("component").
          setInstanceId("instance-id").
          setInstanceIndex(1).
          build();
      return Metrics.MetricPublisherRegisterRequest.newBuilder().setPublisher(publisher).build();
    }

    @Override
    public Message.Builder getResponseBuilder() {
      return Metrics.MetricPublisherRegisterResponse.newBuilder();
    }
  }

  private class MetricsManagerClientResponseHandler
      implements HeronServerTester.TestResponseHandler {
    private int maxMessages;

    MetricsManagerClientResponseHandler(int maxMessages) {
      this.maxMessages = maxMessages;
    }

    @Override
    public void handleResponse(HeronClient client, StatusCode status,
                               Object ctx, Message response) {
      for (int i = 0; i < maxMessages; i++) {
        sendMessage(client);
      }
    }

    private void sendMessage(HeronClient client) {
      Metrics.MetricPublisherPublishMessage.Builder builder =
          Metrics.MetricPublisherPublishMessage.newBuilder();

      for (int j = 0; j < METRICS_COUNT; j++) {
        builder.addMetrics(
            Metrics.MetricDatum.newBuilder()
                    .setName(METRIC_NAME)
                    .setValue(METRIC_VALUE).build());
      }

      for (int j = 0; j < METRICS_COUNT; j++) {
        builder.addExceptions(
            Metrics.ExceptionData.newBuilder()
                .setStacktrace(STACK_TRACE)
                .setLasttime(LAST_TIME)
                .setFirsttime(FIRST_TIME)
                .setCount(EXCEPTION_COUNT)
                .setLogging(LOGGING).build());
      }
      client.sendMessage(builder.build());
    }
  }
}