thread-resource-mgr.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_RUNTIME_THREAD_RESOURCE_MGR_H
#define IMPALA_RUNTIME_THREAD_RESOURCE_MGR_H

#include <stdlib.h>

#include <boost/function.hpp>
#include <boost/scoped_ptr.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/thread/mutex.hpp>
#include <boost/thread/thread.hpp>

#include <list>

#include "common/status.h"

namespace impala {

//
class ThreadResourceMgr {
 public:
  class ResourcePool;

  typedef boost::function<void (ResourcePool*)> ThreadAvailableCb;

  class ResourcePool {
   public:
    void AcquireThreadToken();

    bool TryAcquireThreadToken(bool* is_reserved = NULL);

    void ReserveOptionalTokens(int num);

    void ReleaseThreadToken(bool required);

    void SetThreadAvailableCb(ThreadAvailableCb fn);

    int num_required_threads() const { return num_threads_ & 0xFFFFFFFF; }

    int num_optional_threads() const { return num_threads_ >> 32; }

    int64_t num_threads() const {
      return num_required_threads() + num_optional_threads();
    }

    int num_reserved_optional_threads() { return num_reserved_optional_threads_; }

    bool optional_exceeded() {
      // Cache this so optional/required are computed based on the same value.
      volatile int64_t num_threads = num_threads_;
      int64_t optional_threads = num_threads >> 32;
      int64_t required_threads = num_threads & 0xFFFFFFFF;
      return optional_threads > num_reserved_optional_threads_ &&
             optional_threads + required_threads > quota();
    }

    int num_available_threads() const {
      int value = std::max(quota() - static_cast<int>(num_threads()),
          num_reserved_optional_threads_ - num_optional_threads());
      return std::max(0, value);
    }

    int quota() const { return std::min(max_quota_, parent_->per_pool_quota_); }

    void set_max_quota(int quota) { max_quota_ = quota; }

   private:
    friend class ThreadResourceMgr;

    ResourcePool(ThreadResourceMgr* parent);

    void Reset();

    ThreadResourceMgr* parent_;

    int max_quota_;
    int num_reserved_optional_threads_;

    int64_t num_threads_;

    boost::mutex lock_;

    ThreadAvailableCb thread_available_fn_;
  };

  ThreadResourceMgr(int threads_quota = 0);

  int system_threads_quota() const { return system_threads_quota_; }

  ResourcePool* RegisterPool();

  void UnregisterPool(ResourcePool* pool);

 private:
  int system_threads_quota_;

  boost::mutex lock_;

  typedef std::set<ResourcePool*> Pools;
  Pools pools_;

  int per_pool_quota_;

  std::list<ResourcePool*> free_pool_objs_;

  void UpdatePoolQuotas(ResourcePool* new_pool = NULL);
};

inline void ThreadResourceMgr::ResourcePool::AcquireThreadToken() {
  __sync_fetch_and_add(&num_threads_, 1);
}

inline bool ThreadResourceMgr::ResourcePool::TryAcquireThreadToken(bool* is_reserved) {
  while (true) {
    int64_t previous_num_threads = num_threads_;
    int64_t new_optional_threads = (previous_num_threads >> 32) + 1;
    int64_t new_required_threads = previous_num_threads & 0xFFFFFFFF;
    if (new_optional_threads > num_reserved_optional_threads_ &&
        new_optional_threads + new_required_threads > quota()) {
      return false;
    }
    bool thread_is_reserved = new_optional_threads <= num_reserved_optional_threads_;
    int64_t new_value = new_optional_threads << 32 | new_required_threads;
    // Atomically swap the new value if no one updated num_threads_.  We do not
    // not care about the ABA problem here.
    if (__sync_bool_compare_and_swap(&num_threads_, previous_num_threads, new_value)) {
      if (is_reserved != NULL) *is_reserved = thread_is_reserved;
      return true;
    }
  }
}

inline void ThreadResourceMgr::ResourcePool::ReleaseThreadToken(bool required) {
  if (required) {
    DCHECK_GT(num_required_threads(), 0);
    __sync_fetch_and_add(&num_threads_, -1);
  } else {
    DCHECK_GT(num_optional_threads(), 0);
    while (true) {
      int64_t previous_num_threads = num_threads_;
      int64_t new_optional_threads = (previous_num_threads >> 32) - 1;
      int64_t new_required_threads = previous_num_threads & 0xFFFFFFFF;
      int64_t new_value = new_optional_threads << 32 | new_required_threads;
      if (__sync_bool_compare_and_swap(&num_threads_, previous_num_threads, new_value)) {
        break;
      }
    }
  }

  if (num_available_threads() > 0 && thread_available_fn_ != NULL) {
    boost::unique_lock<boost::mutex> l(lock_);
    if (num_available_threads() > 0 && thread_available_fn_ != NULL) {
      thread_available_fn_(this);
    }
  }
}

} // namespace impala

#endif