io.druid.query.MetricsEmittingQueryRunner Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of druid-processing Show documentation
A module that is everything required to understands Druid Segments
There is a newer version: 0.12.3
/*
 * Licensed to Metamarkets Group Inc. (Metamarkets) under one
 * or more contributor license agreements. See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership. Metamarkets 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 io.druid.query;

import com.google.common.base.Function;
import com.google.common.base.Strings;
import com.google.common.collect.Maps;
import com.metamx.common.guava.Accumulator;
import com.metamx.common.guava.Sequence;
import com.metamx.common.guava.Yielder;
import com.metamx.common.guava.YieldingAccumulator;
import com.metamx.emitter.service.ServiceEmitter;
import com.metamx.emitter.service.ServiceMetricEvent;

import java.io.IOException;
import java.util.Map;

/**
 */
public class MetricsEmittingQueryRunner implements QueryRunner
{
  private static final String DEFAULT_METRIC_NAME = "query/partial/time";

  private final ServiceEmitter emitter;
  private final Function, ServiceMetricEvent.Builder> builderFn;
  private final QueryRunner queryRunner;
  private final long creationTime;
  private final String metricName;
  private final Map userDimensions;

  public MetricsEmittingQueryRunner(
      ServiceEmitter emitter,
      Function, ServiceMetricEvent.Builder> builderFn,
      QueryRunner queryRunner
  )
  {
    this(emitter, builderFn, queryRunner, DEFAULT_METRIC_NAME, Maps.newHashMap());
  }

  public MetricsEmittingQueryRunner(
      ServiceEmitter emitter,
      Function, ServiceMetricEvent.Builder> builderFn,
      QueryRunner queryRunner,
      long creationTime,
      String metricName,
      Map userDimensions
  )
  {
    this.emitter = emitter;
    this.builderFn = builderFn;
    this.queryRunner = queryRunner;
    this.creationTime = creationTime;
    this.metricName = metricName;
    this.userDimensions = userDimensions;
  }

  public MetricsEmittingQueryRunner(
      ServiceEmitter emitter,
      Function, ServiceMetricEvent.Builder> builderFn,
      QueryRunner queryRunner,
      String metricName,
      Map userDimensions
  )
  {
    this(emitter, builderFn, queryRunner, -1, metricName, userDimensions);
  }


  public MetricsEmittingQueryRunner withWaitMeasuredFromNow()
  {
    return new MetricsEmittingQueryRunner(
        emitter,
        builderFn,
        queryRunner,
        System.currentTimeMillis(),
        metricName,
        userDimensions
    );
  }

  @Override
  public Sequence run(final Query query, final Map responseContext)
  {
    final ServiceMetricEvent.Builder builder = builderFn.apply(query);

    for (Map.Entry userDimension : userDimensions.entrySet()) {
      builder.setDimension(userDimension.getKey(), userDimension.getValue());
    }

    builder.setDimension(DruidMetrics.ID, Strings.nullToEmpty(query.getId()));

    return new Sequence()
    {
      @Override
      public  OutType accumulate(OutType outType, Accumulator accumulator)
      {
        OutType retVal;

        long startTime = System.currentTimeMillis();
        try {
          retVal = queryRunner.run(query, responseContext).accumulate(outType, accumulator);
        }
        catch (RuntimeException e) {
          builder.setDimension(DruidMetrics.STATUS, "failed");
          throw e;
        }
        catch (Error e) {
          builder.setDimension(DruidMetrics.STATUS, "failed");
          throw e;
        }
        finally {
          long timeTaken = System.currentTimeMillis() - startTime;

          emitter.emit(builder.build(metricName, timeTaken));

          if (creationTime > 0) {
            emitter.emit(builder.build("query/wait/time", startTime - creationTime));
          }
        }

        return retVal;
      }

      @Override
      public  Yielder toYielder(OutType initValue, YieldingAccumulator accumulator)
      {
        Yielder retVal;

        long startTime = System.currentTimeMillis();
        try {
          retVal = queryRunner.run(query, responseContext).toYielder(initValue, accumulator);
        }
        catch (RuntimeException e) {
          builder.setDimension(DruidMetrics.STATUS, "failed");
          throw e;
        }
        catch (Error e) {
          builder.setDimension(DruidMetrics.STATUS, "failed");
          throw e;
        }

        return makeYielder(startTime, retVal, builder);
      }

      private  Yielder makeYielder(
          final long startTime,
          final Yielder yielder,
          final ServiceMetricEvent.Builder builder
      )
      {
        return new Yielder()
        {
          @Override
          public OutType get()
          {
            return yielder.get();
          }

          @Override
          public Yielder next(OutType initValue)
          {
            try {
              return makeYielder(startTime, yielder.next(initValue), builder);
            }
            catch (RuntimeException e) {
              builder.setDimension(DruidMetrics.STATUS, "failed");
              throw e;
            }
            catch (Error e) {
              builder.setDimension(DruidMetrics.STATUS, "failed");
              throw e;
            }
          }

          @Override
          public boolean isDone()
          {
            return yielder.isDone();
          }

          @Override
          public void close() throws IOException
          {
            try {
              if (!isDone() && builder.getDimension(DruidMetrics.STATUS) == null) {
                builder.setDimension(DruidMetrics.STATUS, "short");
              }

              long timeTaken = System.currentTimeMillis() - startTime;
              emitter.emit(builder.build(metricName, timeTaken));

              if (creationTime > 0) {
                emitter.emit(builder.build("query/wait/time", startTime - creationTime));
              }
            }
            finally {
              yielder.close();
            }
          }
        };
      }
    };
  }
}