text: /*
   * Copyright (c) 2013 Juniper Networks, Inc. All rights reserved.
   */
  
  #include "db/db_table_walker.h"
  
  #include <tbb/atomic.h>
  
  #include "base/logging.h"
  #include "base/task.h"
  #include "db/db.h"
  #include "db/db_partition.h"
  #include "db/db_table.h"
  #include "db/db_table_partition.h"
  
  int DBTableWalker::walker_task_id_ = -1;
  int DBTableWalker::max_iteration_to_yield_ = kIterationToYield;
  
  class DBTableWalker::Walker {
  public:
      Walker(WalkId id, DBTableWalker *wkmgr, DBTable *table,
             const DBRequestKey *key, WalkFn walker, 
             WalkCompleteFn walk_done);
  
      void StopWalk() {
          should_stop_.fetch_and_store(true);
      }
  
      WalkId  id_;
  
      // Parent walker manager
      DBTableWalker *wkmgr_;
  
      // Table on which walk is done
      DBTable *table_;
  
      // Take the ownership of key passed
      std::auto_ptr<DBRequestKey> key_start_;
  
      WalkFn walker_fn_;
      WalkCompleteFn done_fn_;
  
      // Will be true if Table walk is cancelled
      tbb::atomic<bool> should_stop_;
  
      // check whether iteraton is completed on all Table Partition
      tbb::atomic<long> status_;
  };
  
  class DBTableWalker::Worker : public Task {
  public:
      Worker(Walker *walker, int db_partition_id, const DBRequestKey *key) 
          : Task(walker_task_id_, db_partition_id), walker_(walker), 
            key_start_(key) {
          tbl_partition_ = static_cast<DBTablePartition *>(
              walker_->table_->GetTablePartition(db_partition_id));
      }
  
      virtual bool Run();
      std::string Description() const { return "DBTableWalker::Worker"; }
  
  private:
      DBTableWalker::Walker *walker_;
  
      // Store the last visited node to continue walk
      std::auto_ptr<DBRequestKey> walk_ctx_;
  
      // This is where the walk started
      const DBRequestKey *key_start_;
  
      // Table partition for which this worker was created
      DBTablePartition *tbl_partition_;
  };
  
  static void db_walker_wait() {
      static int walk_sleep_usecs_;
      static bool once;
  
      if (!once) {
          once = true;
  
          char *wait = getenv("DB_WALKER_WAIT_USECS");
          if (wait) walk_sleep_usecs_ = strtoul(wait, NULL, 0);
      }
  
      if (walk_sleep_usecs_) {
          usleep(walk_sleep_usecs_);
      }
  }
  
  bool DBTableWalker::Worker::Run() {
      int count = 0;
      DBRequestKey *key_resume;
  
      // Check whether Walker was requested to be cancelled
      if (walker_->should_stop_) {
          goto walk_done;
      }
  
      // Check where we left in last iteration
      if ((key_resume = walk_ctx_.get()) == NULL) {
          // First time invoke of worker thread, start from key_start_
          key_resume = const_cast <DBRequestKey *>(key_start_);
      }
  
      DBEntry *entry;
      if (key_resume != NULL) {
          DBTable *table = walker_->table_;
          std::auto_ptr<const DBEntryBase> start;
          start = table->AllocEntry(key_resume);
          // Find matching or next in sort order
          entry = tbl_partition_->lower_bound(start.get());
      } else {
          entry = tbl_partition_->GetFirst();
      }
      if (entry == NULL) {
          goto walk_done;
      }
  
      for (DBEntry *next = NULL; entry; entry = next) {
          next = tbl_partition_->GetNext(entry);
          // Check whether Walker was requested to be cancelled
          if (walker_->should_stop_) {
              break; 
          }
          if (count == GetIterationToYield()) {
              // store the context
              walk_ctx_ = entry->GetDBRequestKey();
              return false;
          }
  
          // Invoke walker function
          bool more = walker_->walker_fn_(tbl_partition_, entry);
          if (!more) {
              break;
          }
  
          db_walker_wait();
          count++;
      }
  
  walk_done:
      // Check whether all other walks on the table is completed
      long num_walkers_on_tpart = walker_->status_.fetch_and_decrement();
      if (num_walkers_on_tpart == 1) {
          if (walker_->should_stop_) {
              walker_->table_->incr_walk_cancel_count();
          } else {
              walker_->table_->incr_walk_complete_count();
              // Invoke Walker_Complete callback
              if (walker_->done_fn_ != NULL) {
                  walker_->done_fn_(walker_->table_);
              }
          }
  
          // Release the memory for walker and bitmap
          walker_->wkmgr_->PurgeWalker(walker_->id_);
      }
      return true;
  }
  
  DBTableWalker::Walker::Walker(WalkId id, DBTableWalker *wkmgr,
                                DBTable *table, const DBRequestKey *key,
                                WalkFn walker, WalkCompleteFn walk_done)
      : id_(id), wkmgr_(wkmgr), table_(table),
        key_start_(const_cast<DBRequestKey *>(key)), 
        walker_fn_(walker), done_fn_(walk_done) {
      int num_worker = table->PartitionCount();
      should_stop_ = false;
      status_ = num_worker;
      for (int i = 0; i < num_worker; i++) {
          Worker *task = new Worker(this, i, key);
          TaskScheduler *scheduler = TaskScheduler::GetInstance();
          scheduler->Enqueue(task);
      }
  }
  
  DBTableWalker::DBTableWalker() {
      if (walker_task_id_ == -1) {
          TaskScheduler *scheduler = TaskScheduler::GetInstance();
          // Using same task id as DBPartition
          walker_task_id_ = scheduler->GetTaskId("db::DBTable");
      }
  }
  
  DBTableWalker::WalkId DBTableWalker::WalkTable(DBTable *table, 
                                                 const DBRequestKey *key_start, 
                                                 WalkFn walkerfn , 
                                                 WalkCompleteFn walk_complete) {
      table->incr_walk_request_count();
      tbb::mutex::scoped_lock lock(walkers_mutex_);
      size_t i = walker_map_.find_first();
      if (i == walker_map_.npos) {
          i = walkers_.size();
          Walker *walker = new Walker(i, this, table, key_start, 
                                      walkerfn, walk_complete);
          walkers_.push_back(walker);
      } else {
          walker_map_.reset(i);
          if (walker_map_.none()) {
              walker_map_.clear();
          }
          Walker *walker = new Walker(i, this, table, key_start, 
                                      walkerfn, walk_complete);
          walkers_[i] = walker;
      }
      table->incr_walker_count();
      return i;
  }
  
  void DBTableWalker::WalkCancel(WalkId id) {
      tbb::mutex::scoped_lock lock(walkers_mutex_);
      walkers_[id]->StopWalk();
      // Purge to be called after task has stopped
  }
  
  void DBTableWalker::PurgeWalker(WalkId id) {
      tbb::mutex::scoped_lock lock(walkers_mutex_);
      Walker *walker = walkers_[id];
      DBTable *table = walker->table_;
      delete walker;
      walkers_[id] = NULL;
      if ((size_t) id == walkers_.size() - 1) {
          while (!walkers_.empty() && walkers_.back() == NULL) {
              walkers_.pop_back();
          }
          if (walker_map_.size() > walkers_.size()) {
              walker_map_.resize(walkers_.size());
          }
      } else {
          if ((size_t) id >= walker_map_.size()) {
              walker_map_.resize(id + 1);
          }
          walker_map_.set(id);
      }
  
      // Retry table deletion when the last walker is purged.
      if (table->decr_walker_count() == 0) {
          table->RetryDelete();
      }
  }
  