Function UpdateStats

tensorflow/contrib/tensor_forest/kernels/stats_ops.cc:148–195

Try to update a leaf's stats by acquiring its lock. If it can't be acquired, put it in a waiting queue to come back to later and try the next one. Once all leaf_ids have been visited, cycle through the waiting ids until they're gone.

// one.  Once all leaf_ids have been visited, cycle through the waiting ids
// until they're gone.
void UpdateStats(const core::RefCountPtr<FertileStatsResource>& resource,
                 const std::unique_ptr<TensorDataSet>& data,
                 const TensorInputTarget& target, int num_targets,
                 const Tensor& leaf_ids_tensor,
                 std::unordered_map<int32, std::unique_ptr<mutex>>* locks,
                 mutex* set_lock, int32 start, int32 end,
                 std::unordered_set<int32>* ready_to_split) {
  const auto leaf_ids = leaf_ids_tensor.unaligned_flat<int32>();

  // Stores leaf_id, leaf_depth, example_id for examples that are waiting
  // on another to finish.
  std::queue<std::tuple<int32, int32>> waiting;

  int32 i = start;
  while (i < end || !waiting.empty()) {
    int32 leaf_id;
    int32 example_id;
    bool was_waiting = false;
    if (i >= end) {
      std::tie(leaf_id, example_id) = waiting.front();
      waiting.pop();
      was_waiting = true;
    } else {
      leaf_id = leaf_ids(i);
      example_id = i;
      ++i;
    }
    const std::unique_ptr<mutex>& leaf_lock = (*locks)[leaf_id];
    if (was_waiting) {
      leaf_lock->lock();
    } else {
      if (!leaf_lock->try_lock()) {
        waiting.emplace(leaf_id, example_id);
        continue;
      }
    }

    bool is_finished;
    resource->AddExampleToStatsAndInitialize(data, &target, {example_id},
                                             leaf_id, &is_finished);
    leaf_lock->unlock();
    if (is_finished) {
      set_lock->lock();
      ready_to_split->insert(leaf_id);
      set_lock->unlock();
    }
  }
}

// Update leaves from start through end in the leaf_examples iterator.
void UpdateStatsCollated(