data-stream-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_DATA_STREAM_MGR_H
#define IMPALA_RUNTIME_DATA_STREAM_MGR_H

#include <list>
#include <set>
#include <boost/thread/mutex.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/thread/condition_variable.hpp>
#include <boost/unordered_map.hpp>

#include "common/status.h"
#include "common/object-pool.h"
#include "runtime/descriptors.h"  // for PlanNodeId
#include "runtime/mem-tracker.h"
#include "util/promise.h"
#include "util/runtime-profile.h"
#include "gen-cpp/Types_types.h"  // for TUniqueId

namespace impala {

class DescriptorTbl;
class DataStreamRecvr;
class RowBatch;
class RuntimeState;
class TRowBatch;

//
//
class DataStreamMgr {
 public:
  DataStreamMgr() {}

  boost::shared_ptr<DataStreamRecvr> CreateRecvr(
      RuntimeState* state, const RowDescriptor& row_desc,
      const TUniqueId& fragment_instance_id, PlanNodeId dest_node_id,
      int num_senders, int buffer_size, RuntimeProfile* profile,
      bool is_merging);

  Status AddData(const TUniqueId& fragment_instance_id, PlanNodeId dest_node_id,
                 const TRowBatch& thrift_batch, int sender_id);

  Status CloseSender(const TUniqueId& fragment_instance_id, PlanNodeId dest_node_id,
      int sender_id);

  void Cancel(const TUniqueId& fragment_instance_id);

 private:
  friend class DataStreamRecvr;

  boost::mutex lock_;

  typedef boost::unordered_multimap<uint32_t,
      boost::shared_ptr<DataStreamRecvr> > StreamMap;
  StreamMap receiver_map_;

  struct ComparisonOp {
    bool operator()(const std::pair<impala::TUniqueId, PlanNodeId>& a,
                    const std::pair<impala::TUniqueId, PlanNodeId>& b) {
      if (a.first.hi < b.first.hi) {
        return true;
      } else if (a.first.hi > b.first.hi) {
        return false;
      } else if (a.first.lo < b.first.lo) {
        return true;
      } else if (a.first.lo > b.first.lo) {
        return false;
      }
      return a.second < b.second;
    }
  };

  typedef std::set<std::pair<TUniqueId, PlanNodeId>, ComparisonOp > FragmentStreamSet;
  FragmentStreamSet fragment_stream_set_;

  boost::shared_ptr<DataStreamRecvr> FindRecvr(
      const TUniqueId& fragment_instance_id, PlanNodeId node_id,
      bool acquire_lock = true);

  Status DeregisterRecvr(const TUniqueId& fragment_instance_id, PlanNodeId node_id);

  inline uint32_t GetHashValue(const TUniqueId& fragment_instance_id, PlanNodeId node_id);
};

}

#endif