streaming-sampler.h

Go to the documentation of this file.

// Copyright 2012 Cloudera 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.

#ifndef IMPALA_UTIL_STREAMING_SAMPLER_H
#define IMPALA_UTIL_STREAMING_SAMPLER_H

#include <string.h>
#include <iostream>
#include "util/spinlock.h"

namespace impala {

template<typename T, int MAX_SAMPLES>
class StreamingSampler {
 public:
  StreamingSampler(int initial_period = 500)
    : samples_collected_(0) ,
      period_(initial_period),
      current_sample_sum_(0),
      current_sample_count_(0),
      current_sample_total_time_(0) {
  }

  StreamingSampler(int period, const std::vector<T>& initial_samples)
    : samples_collected_(initial_samples.size()),
      period_(period),
      current_sample_sum_(0),
      current_sample_count_(0),
      current_sample_total_time_(0) {
    DCHECK_LE(samples_collected_, MAX_SAMPLES);
    memcpy(samples_, &initial_samples[0], sizeof(T) * samples_collected_);
  }

  void AddSample(T sample, int ms) {
    boost::lock_guard<SpinLock> l(lock_);
    ++current_sample_count_;
    current_sample_sum_ += sample;
    current_sample_total_time_ += ms;

    if (current_sample_total_time_ >= period_) {
      samples_[samples_collected_++] = current_sample_sum_ / current_sample_count_;
      current_sample_count_ = 0;
      current_sample_sum_ = 0;
      current_sample_total_time_ = 0;

      if (samples_collected_ == MAX_SAMPLES) {
        period_ *= 2;
        for (int i = 0; i < MAX_SAMPLES / 2; ++i) {
          samples_[i] = (samples_[i * 2] + samples_[i * 2 + 1]) / 2;
        }
        samples_collected_ /= 2;
      }
    }
  }

  const T* GetSamples(int* num_samples, int* period, SpinLock** lock = NULL) const {
    if (lock != NULL) {
      lock_.lock();
      *lock = &lock_;
    }
    *num_samples = samples_collected_;
    *period = period_;
    return samples_;
  }

  void SetSamples(int period, const std::vector<T>& samples) {
    DCHECK_LE(samples.size(), MAX_SAMPLES);

    boost::lock_guard<SpinLock> l(lock_);
    period_ = period;
    samples_collected_ = samples.size();
    memcpy(samples_, &samples[0], sizeof(T) * samples_collected_);
    current_sample_sum_ = 0;
    current_sample_count_ = 0;
    current_sample_total_time_ = 0;
  }

  std::string DebugString(const std::string& prefix="") const {
    boost::lock_guard<SpinLock> l(lock_);
    std::stringstream ss;
    ss << prefix << "Period = " << period_ << std::endl
       << prefix << "Num = " << samples_collected_ << std::endl
       << prefix << "Samples = {";
    for (int i = 0; i < samples_collected_; ++i) {
      ss << samples_[i] << ", ";
    }
    ss << prefix << "}" << std::endl;
    return ss.str();
  }

 private:
  mutable SpinLock lock_;

  T samples_[MAX_SAMPLES];

  int samples_collected_;

  int period_;

  T current_sample_sum_;

  int current_sample_count_;

  int current_sample_total_time_;
};

}

#endif