state-store-2.0-test.cc

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.

#include <vector>

#include <boost/date_time/posix_time/posix_time.hpp>
#include <boost/foreach.hpp>
#include <boost/scoped_ptr.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/thread/barrier.hpp>
#include <boost/thread/thread_time.hpp>
#include <boost/unordered_set.hpp>
#include <gtest/gtest.h>

#include "common/daemon.h"
#include "common/status.h"
#include "statestore/state-store-subscriber-2.0.h"
#include "testutil/in-process-servers.h"
#include "util/cpu-info.h"
#include "util/metrics.h"
#include "util/network-util.h"
#include "util/thrift-util.h"
#include "gen-cpp/StatestoreTypes_types.h"
#include "gen-cpp/Types_types.h"

using namespace apache::thrift;
using namespace apache::thrift::transport;
using namespace boost;
using namespace std;

DECLARE_int32(rpc_cnxn_attempts);
DECLARE_int32(rpc_cnxn_retry_interval_ms);
DECLARE_int32(statestore_max_missed_heartbeats);

namespace impala {

class StateStore2Test : public testing::Test {
 public:
  class Callbacks {
   public:
    virtual void Topic1(const StateStoreSubscriber2::TopicDeltaMap& state,
                        vector<TTopicUpdate>* topic_updates) {
    }

    virtual void Topic2(const StateStoreSubscriber2::TopicDeltaMap& state,
                        vector<TTopicUpdate>* topic_updates) {
    }
  };

 protected:
  class TopicCallbacks {
   public:
    TopicCallbacks(Callbacks* cb) : callbacks_(cb) { }

    void Topic1(const StateStoreSubscriber2::TopicDeltaMap& state,
                vector<TTopicUpdate>* topic_updates) {
      callbacks_->Topic1(state, topic_updates);
    }

    void Topic2(const StateStoreSubscriber2::TopicDeltaMap& state,
                vector<TTopicUpdate>* topic_updates) {
      callbacks_->Topic2(state, topic_updates);
    }

    StateStoreSubscriber2::UpdateCallback topic1_callback() {
      return bind<void>(mem_fn(&TopicCallbacks::Topic1), this, _1, _2);
    }

    StateStoreSubscriber2::UpdateCallback topic2_callback() {
      return bind<void>(mem_fn(&TopicCallbacks::Topic2), this, _1, _2);
    }

    ~TopicCallbacks() {
      delete callbacks_;
    }

   private:
    Callbacks* callbacks_;
  };

  static int next_port_;
  scoped_ptr<InProcessStateStore> state_store_;
  static const char* ipaddress_;


  vector<shared_ptr<StateStoreSubscriber2> > subscribers_;
  vector<TopicCallbacks*> callbacks_;

  StateStore2Test()
      : state_store_(new InProcessStateStore(next_port_++, next_port_++)) {
  }

  virtual void SetUp() {
    LOG(INFO) << "SETTING UP";

    Status status = state_store_->Start();
    VLOG(1) << status.GetErrorMsg();
    EXPECT_TRUE(status.ok());
  }

  virtual void TearDown() {
    EXPECT_TRUE(state_store_->StopForTesting().ok());
    // TODO: delete all the callbacks
  }

  shared_ptr<StateStoreSubscriber2> StartStateStoreSubscriber(Callbacks* cb) {
    TNetworkAddress statestore = MakeNetworkAddress("localhost", state_store_->port());
    TNetworkAddress subscriber = MakeNetworkAddress("localhost", next_port_++);
    stringstream subscriber_id;
    subscriber_id << subscriber;
    subscribers_.push_back(shared_ptr<StateStoreSubscriber2>(
        new StateStoreSubscriber2(subscriber_id.str(), subscriber, statestore)));

    callbacks_.push_back(new TopicCallbacks(cb));
    subscribers_.back()->AddTopic("topic_1", false, callbacks_.back()->topic1_callback());
    subscribers_.back()->AddTopic("topic_2", true, callbacks_.back()->topic2_callback());

    EXPECT_TRUE(subscribers_.back()->Start().ok());
    EXPECT_TRUE(WaitForServer("localhost", subscriber.port, 10, 100).ok());
    return subscribers_.back();
  }

  void UnregisterAllSubscribers() {
    // TODO: Stop for testing?
  }
};

int StateStore2Test::next_port_ = 27000;

Status WaitForCV(unique_lock<mutex>* lock, condition_variable* cv) {
  system_time deadline = get_system_time() +
      posix_time::milliseconds(5000);

  if (!cv->timed_wait(*lock, deadline)) {
    return Status("CV deadline expired");
  }

  return Status::OK;
}


class TestRegistrationCb : public StateStore2Test::Callbacks {
 public:
  virtual void Topic1(const StateStoreSubscriber2::TopicDeltaMap& state,
                      vector<TTopicUpdate>* topic_updates) {
    cv_.notify_all();
  }

  mutex* lock1() { return &cv_mutex_; }
  condition_variable* cv1() { return &cv_; }

 protected:
  mutex cv_mutex_;
  condition_variable cv_;

};

TEST_F(StateStore2Test, TestRegistration) {
  TestRegistrationCb* cb = new TestRegistrationCb();
  unique_lock<mutex> l(*cb->lock1());
  shared_ptr<StateStoreSubscriber2> running_subscriber =
      StartStateStoreSubscriber(cb);
  EXPECT_TRUE(WaitForCV(&l, cb->cv1()).ok());
}

class TestReceiveTopicCb : public TestRegistrationCb {
  virtual void Topic1(const StateStoreSubscriber2::TopicDeltaMap& state,
                      vector<TTopicUpdate>* topic_updates) {
    StateStoreSubscriber2::TopicDeltaMap::const_iterator topic =
        state.find("topic_1");
    if (topic == state.end() || topic->second.topic_entries.size() == 0) {
      topic_updates->push_back(TTopicUpdate());
      TTopicUpdate& update = topic_updates->back();
      update.topic_name = "topic_1";
      update.topic_updates.push_back(TTopicItem());
      TTopicItem& item = update.topic_updates.back();
      item.key = "hello world";
      item.value = "value";
    } else {
      const TTopicDelta& delta = topic->second;
      if (delta.topic_entries.size() > 0 && delta.topic_entries[0].key == "hello world") {
        cv_.notify_all();
      }
    }
  }
};

TEST_F(StateStore2Test, TestReceiveUpdates) {
  TestReceiveTopicCb* cb = new TestReceiveTopicCb();
  unique_lock<mutex> l(*cb->lock1());
  shared_ptr<StateStoreSubscriber2> running_subscriber =
      StartStateStoreSubscriber(cb);
  EXPECT_TRUE(WaitForCV(&l, cb->cv1()).ok());
}

}

using namespace impala;

int main(int argc, char **argv) {
  InitDaemon(argc, argv);

  ::testing::InitGoogleTest(&argc, argv);
  return RUN_ALL_TESTS();
}