mem-tracker.h

Go to the documentation of this file.

// Copyright 2013 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_RUNTIME_MEM_TRACKER_H
#define IMPALA_RUNTIME_MEM_TRACKER_H

#include <stdint.h>
#include <map>
#include <vector>
#include <boost/shared_ptr.hpp>
#include <boost/weak_ptr.hpp>
#include <boost/thread/mutex.hpp>
#include <boost/unordered_map.hpp>

#include "common/logging.h"
#include "common/atomic.h"
#include "util/debug-util.h"
#include "util/internal-queue.h"
#include "util/metrics.h"
#include "util/runtime-profile.h"
#include "util/spinlock.h"

#include "gen-cpp/Types_types.h" // for TUniqueId

namespace impala {

class MemTracker;
class QueryResourceMgr;

//
//
//
class MemTracker {
 public:
  MemTracker(int64_t byte_limit = -1, int64_t rm_reserved_limit = -1,
      const std::string& label = std::string(), MemTracker* parent = NULL,
      bool log_usage_if_zero = true);

  MemTracker(RuntimeProfile* profile, int64_t byte_limit, int64_t rm_reserved_limit = -1,
      const std::string& label = std::string(), MemTracker* parent = NULL);

  MemTracker(UIntGauge* consumption_metric,
      int64_t byte_limit = -1, int64_t rm_reserved_limit = -1,
      const std::string& label = std::string());

  ~MemTracker();

  void UnregisterFromParent();

  static boost::shared_ptr<MemTracker> GetQueryMemTracker(const TUniqueId& id,
      int64_t byte_limit, int64_t rm_reserved_limit, MemTracker* parent,
      QueryResourceMgr* res_mgr);

  static MemTracker* GetRequestPoolMemTracker(const std::string& pool_name,
      MemTracker* parent);

  int64_t effective_limit() const {
    // TODO: maybe no limit should be MAX_LONG?
    DCHECK(rm_reserved_limit_ <= limit_ || limit_ == -1);
    if (rm_reserved_limit_ == -1) return limit_;
    return rm_reserved_limit_;
  }

  void Consume(int64_t bytes) {
    if (bytes < 0) {
      Release(-bytes);
      return;
    }

    if (consumption_metric_ != NULL) {
      DCHECK(parent_ == NULL);
      consumption_->Set(consumption_metric_->value());
      return;
    }
    if (bytes == 0) return;
    if (UNLIKELY(enable_logging_)) LogUpdate(true, bytes);
    for (std::vector<MemTracker*>::iterator tracker = all_trackers_.begin();
         tracker != all_trackers_.end(); ++tracker) {
      (*tracker)->consumption_->Add(bytes);
      if ((*tracker)->consumption_metric_ == NULL) {
        DCHECK_GE((*tracker)->consumption_->current_value(), 0);
      }
    }
  }

  void ConsumeLocal(int64_t bytes, MemTracker* end_tracker) {
    DCHECK(consumption_metric_ == NULL) << "Should not be called on root.";
    if (UNLIKELY(enable_logging_)) LogUpdate(bytes > 0, bytes);
    for (int i = 0; i < all_trackers_.size(); ++i) {
      if (all_trackers_[i] == end_tracker) return;
      DCHECK(!all_trackers_[i]->has_limit());
      all_trackers_[i]->consumption_->Add(bytes);
    }
  }

  void ReleaseLocal(int64_t bytes, MemTracker* end_tracker) {
    ConsumeLocal(-bytes, end_tracker);
  }

  bool TryConsume(int64_t bytes) {
    if (consumption_metric_ != NULL) consumption_->Set(consumption_metric_->value());
    if (bytes <= 0) return true;
    if (UNLIKELY(enable_logging_)) LogUpdate(true, bytes);
    int i = 0;
    // Walk the tracker tree top-down, to avoid expanding a limit on a child whose parent
    // won't accommodate the change.
    for (i = all_trackers_.size() - 1; i >= 0; --i) {
      MemTracker* tracker = all_trackers_[i];
      int64_t limit = tracker->effective_limit();
      if (limit < 0) {
        tracker->consumption_->Add(bytes);
      } else {
        if (!tracker->consumption_->TryAdd(bytes, limit)) {
          // One of the trackers failed, attempt to GC memory or expand our limit. If that
          // succeeds, TryUpdate() again. Bail if either fails.
          //
          // TODO: This may not be right if more than one tracker can actually change its
          // rm reservation limit.
          if (!tracker->GcMemory(limit - bytes) || tracker->ExpandRmReservation(bytes)) {
            if (!tracker->consumption_->TryAdd(bytes, tracker->limit_)) break;
          } else {
            break;
          }
        }
      }
    }
    // Everyone succeeded, return.
    if (i == -1) return true;

    // Someone failed, roll back the ones that succeeded.
    // TODO: this doesn't roll it back completely since the max values for
    // the updated trackers aren't decremented. The max values are only used
    // for error reporting so this is probably okay. Rolling those back is
    // pretty hard; we'd need something like 2PC.
    //
    // TODO: This might leave us with an allocated resource that we can't use. Do we need
    // to adjust the consumption of the query tracker to stop the resource from never
    // getting used by a subsequent TryConsume()?
    for (int j = all_trackers_.size() - 1; j > i; --j) {
      all_trackers_[j]->consumption_->Add(-bytes);
    }
    return false;
  }

  void Release(int64_t bytes) {
    if (bytes < 0) {
      Consume(-bytes);
      return;
    }

    if (UNLIKELY(released_memory_since_gc_.UpdateAndFetch(bytes)) > GC_RELEASE_SIZE) {
      GcTcmalloc();
    }

    if (consumption_metric_ != NULL) {
      DCHECK(parent_ == NULL);
      consumption_->Set(consumption_metric_->value());
      return;
    }
    if (bytes == 0) return;
    if (UNLIKELY(enable_logging_)) LogUpdate(false, bytes);
    for (std::vector<MemTracker*>::iterator tracker = all_trackers_.begin();
         tracker != all_trackers_.end(); ++tracker) {
      (*tracker)->consumption_->Add(-bytes);
      if ((*tracker)->consumption_metric_ == NULL) {
        DCHECK_GE((*tracker)->consumption_->current_value(), 0)
          << std::endl << (*tracker)->LogUsage();
      }
    }

  }

  bool AnyLimitExceeded() {
    for (std::vector<MemTracker*>::iterator tracker = limit_trackers_.begin();
         tracker != limit_trackers_.end(); ++tracker) {
      if ((*tracker)->LimitExceeded()) return true;
    }
    return false;
  }

  bool LimitExceeded() {
    if (UNLIKELY(CheckLimitExceeded())) {
      if (bytes_over_limit_metric_ != NULL) {
        bytes_over_limit_metric_->set_value(consumption() - limit_);
      }
      return GcMemory(limit_);
    }
    return false;
  }

  int64_t SpareCapacity() const {
    int64_t result = std::numeric_limits<int64_t>::max();
    for (std::vector<MemTracker*>::const_iterator tracker = limit_trackers_.begin();
         tracker != limit_trackers_.end(); ++tracker) {
      int64_t mem_left = (*tracker)->limit() - (*tracker)->consumption();
      result = std::min(result, mem_left);
    }
    return result;
  }


  int64_t limit() const { return limit_; }
  bool has_limit() const { return limit_ >= 0; }
  const std::string& label() const { return label_; }

  int64_t lowest_limit() const {
    if (limit_trackers_.empty()) return -1;
    int64_t v = std::numeric_limits<int64_t>::max();
    for (int i = 0; i < limit_trackers_.size(); ++i) {
      DCHECK(limit_trackers_[i]->has_limit());
      v = std::min(v, limit_trackers_[i]->limit());
    }
    return v;
  }

  int64_t consumption() const { return consumption_->current_value(); }

  int64_t peak_consumption() const { return consumption_->value(); }

  MemTracker* parent() const { return parent_; }

  typedef boost::function<void ()> GcFunction;

  void AddGcFunction(GcFunction f) { gc_functions_.push_back(f); }

  void RegisterMetrics(MetricGroup* metrics, const std::string& prefix);

  std::string LogUsage(const std::string& prefix = "") const;

  void EnableLogging(bool enable, bool log_stack) {
    enable_logging_ = enable;
    log_stack_ = log_stack;
  }

  static const std::string COUNTER_NAME;

 private:
  bool CheckLimitExceeded() const { return limit_ >= 0 && limit_ < consumption(); }

  bool GcMemory(int64_t max_consumption);

  void GcTcmalloc();

  void SetQueryResourceMgr(QueryResourceMgr* context) {
    query_resource_mgr_ = context;
  }

  void Init();

  void AddChildTracker(MemTracker* tracker);

  void LogUpdate(bool is_consume, int64_t bytes) const;

  static std::string LogUsage(const std::string& prefix,
      const std::list<MemTracker*>& trackers);

  bool ExpandRmReservation(int64_t bytes);

  static const int64_t GC_RELEASE_SIZE = 128 * 1024L * 1024L;

  static AtomicInt<int64_t> released_memory_since_gc_;

  SpinLock gc_lock_;

  static boost::mutex static_mem_trackers_lock_;

  typedef boost::unordered_map<TUniqueId, boost::weak_ptr<MemTracker> >
  RequestTrackersMap;
  static RequestTrackersMap request_to_mem_trackers_;

  typedef boost::unordered_map<std::string, MemTracker*> PoolTrackersMap;
  static PoolTrackersMap pool_to_mem_trackers_;

  TUniqueId query_id_;

  std::string pool_name_;

  int64_t limit_;

  int64_t rm_reserved_limit_;

  std::string label_;
  MemTracker* parent_;

  RuntimeProfile::HighWaterMarkCounter* consumption_;

  RuntimeProfile::HighWaterMarkCounter local_counter_;

  UIntGauge* consumption_metric_;

  std::vector<MemTracker*> all_trackers_;  // this tracker plus all of its ancestors
  std::vector<MemTracker*> limit_trackers_;  // all_trackers_ with valid limits

  mutable boost::mutex child_trackers_lock_;
  std::list<MemTracker*> child_trackers_;

  std::list<MemTracker*>::iterator child_tracker_it_;

  std::vector<GcFunction> gc_functions_;

  bool auto_unregister_;

  bool enable_logging_;
  bool log_stack_;

  bool log_usage_if_zero_;

  boost::mutex resource_acquisition_lock_;

  QueryResourceMgr* query_resource_mgr_;

  IntCounter* num_gcs_metric_;

  IntGauge* bytes_freed_by_last_gc_metric_;

  IntGauge* bytes_over_limit_metric_;
};

}

#endif